'Base Command'
python trainer.py --dataset_path /home/meow/ACE-Step/Zen/dataset.json --checkpoint_dir /home/meow/.cache/ace-step/checkpoints --exp_name "DaltOniX"

'Flags'
--dataset_path: Path to your Huggingface dataset (required)
--checkpoint_dir: Directory containing the base model checkpoint
--learning_rate: Learning rate for training (default: 1e-4)
--max_steps: Maximum number of training steps (default: 2000000)
--precision: Training precision, e.g., "bf16-mixed" (default) or "fp32"
--devices: Number of GPUs to use (default: 1)
--num_nodes: Number of compute nodes to use (default: 1)
--accumulate_grad_batches: Gradient accumulation steps (default: 1)
--num_workers: Number of data loading workers (default: 8)
--every_n_train_steps: Checkpoint saving frequency (default: 2000)
--every_plot_step: Frequency of generating evaluation samples (default: 2000)
--exp_name: Experiment name for logging (default: "text2music_train_test")
--logger_dir: Directory for saving logs (default: "./exps/logs/")

'My Problem'
I'm getting out of memory errors and all kinds of weird crap please construct a command that will train well on 12 GB vram also here is some past error logs...

'Past Error Logs'
(ace_step) meow@Anbernic-Win600:~/ACE-Step$ python trainer.py --dataset_path /home/meow/ACE-Step/Zen/dataset.json --checkpoint_dir /home/meow/.cache/ace-step/checkpoints --exp_name "DaltOniX" --precision bf16-mixed --devices 1 --num_nodes 1 --accumulate_grad_batches 2 --num_workers 4 --learning_rate 1e-4 --max_steps 50000 --every_n_train_steps 1000 --every_plot_step 1000
Fetching 13 files: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 13/13 [00:00<00:00, 244511.00it/s]
[2025-05-10 20:17:19,054] [INFO] [real_accelerator.py:239:get_accelerator] Setting ds_accelerator to cuda (auto detect)
GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
initializing deepspeed distributed: GLOBAL_RANK: 0, MEMBER: 1/1
[rank0]: Traceback (most recent call last):
[rank0]:   File "/home/meow/ACE-Step/trainer.py", line 853, in <module>
[rank0]:     main(args)
[rank0]:   File "/home/meow/ACE-Step/trainer.py", line 823, in main
[rank0]:     trainer.fit(
[rank0]:   File "/home/meow/anaconda3/envs/ace_step/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 561, in fit
[rank0]:     call._call_and_handle_interrupt(
[rank0]:   File "/home/meow/anaconda3/envs/ace_step/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 47, in _call_and_handle_interrupt
[rank0]:     return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
[rank0]:   File "/home/meow/anaconda3/envs/ace_step/lib/python3.10/site-packages/pytorch_lightning/strategies/launchers/subprocess_script.py", line 105, in launch
[rank0]:     return function(*args, **kwargs)
[rank0]:   File "/home/meow/anaconda3/envs/ace_step/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 599, in _fit_impl
[rank0]:     self._run(model, ckpt_path=ckpt_path)
[rank0]:   File "/home/meow/anaconda3/envs/ace_step/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 988, in _run
[rank0]:     self.strategy.setup(self)
[rank0]:   File "/home/meow/anaconda3/envs/ace_step/lib/python3.10/site-packages/pytorch_lightning/strategies/deepspeed.py", line 341, in setup
[rank0]:     self._init_config_if_needed()
[rank0]:   File "/home/meow/anaconda3/envs/ace_step/lib/python3.10/site-packages/pytorch_lightning/strategies/deepspeed.py", line 805, in _init_config_if_needed
[rank0]:     self._format_config()
[rank0]:   File "/home/meow/anaconda3/envs/ace_step/lib/python3.10/site-packages/pytorch_lightning/strategies/deepspeed.py", line 814, in _format_config
[rank0]:     self._format_batch_size_and_grad_accum_config()
[rank0]:   File "/home/meow/anaconda3/envs/ace_step/lib/python3.10/site-packages/pytorch_lightning/strategies/deepspeed.py", line 907, in _format_batch_size_and_grad_accum_config
[rank0]:     batch_size = self._auto_select_batch_size()
[rank0]:   File "/home/meow/anaconda3/envs/ace_step/lib/python3.10/site-packages/pytorch_lightning/strategies/deepspeed.py", line 919, in _auto_select_batch_size
[rank0]:     train_dataloader = data_source.dataloader()
[rank0]:   File "/home/meow/anaconda3/envs/ace_step/lib/python3.10/site-packages/pytorch_lightning/trainer/connectors/data_connector.py", line 290, in dataloader
[rank0]:     return call._call_lightning_module_hook(self.instance.trainer, self.name, pl_module=self.instance)
[rank0]:   File "/home/meow/anaconda3/envs/ace_step/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 176, in _call_lightning_module_hook
[rank0]:     output = fn(*args, **kwargs)
[rank0]:   File "/home/meow/ACE-Step/trainer.py", line 422, in train_dataloader
[rank0]:     self.train_dataset = Text2MusicDataset(
[rank0]:   File "/home/meow/ACE-Step/acestep/text2music_dataset.py", line 204, in __init__
[rank0]:     self.setup_full(train, shuffle, sample_size)
[rank0]:   File "/home/meow/ACE-Step/acestep/text2music_dataset.py", line 216, in setup_full
[rank0]:     pretrain_ds = load_from_disk(self.train_dataset_path)
[rank0]:   File "/home/meow/anaconda3/envs/ace_step/lib/python3.10/site-packages/datasets/load.py", line 2147, in load_from_disk
[rank0]:     raise FileNotFoundError(
[rank0]: FileNotFoundError: Directory /home/meow/ACE-Step/Zen/dataset.json is neither a `Dataset` directory nor a `DatasetDict` directory.
[rank0]:[W510 20:18:08.096005805 ProcessGroupNCCL.cpp:1476] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())

'More'
(ace_step) meow@Anbernic-Win600:~/ACE-Step$ python trainer.py --dataset_path /home/meow/ACE-Step/Zen/dataset.json --checkpoint_dir /home/meow/.cache/ace-step/checkpoints --exp_name "DaltOniX"
Fetching 13 files: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 13/13 [00:00<00:00, 5987.26it/s]
[2025-05-10 20:10:15,507] [INFO] [real_accelerator.py:239:get_accelerator] Setting ds_accelerator to cuda (auto detect)
GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
initializing deepspeed distributed: GLOBAL_RANK: 0, MEMBER: 1/1
[rank0]: Traceback (most recent call last):
[rank0]:   File "/home/meow/ACE-Step/trainer.py", line 853, in <module>
[rank0]:     main(args)
[rank0]:   File "/home/meow/ACE-Step/trainer.py", line 823, in main
[rank0]:     trainer.fit(
[rank0]:   File "/home/meow/anaconda3/envs/ace_step/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 561, in fit
[rank0]:     call._call_and_handle_interrupt(
[rank0]:   File "/home/meow/anaconda3/envs/ace_step/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 47, in _call_and_handle_interrupt
[rank0]:     return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
[rank0]:   File "/home/meow/anaconda3/envs/ace_step/lib/python3.10/site-packages/pytorch_lightning/strategies/launchers/subprocess_script.py", line 105, in launch
[rank0]:     return function(*args, **kwargs)
[rank0]:   File "/home/meow/anaconda3/envs/ace_step/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 599, in _fit_impl
[rank0]:     self._run(model, ckpt_path=ckpt_path)
[rank0]:   File "/home/meow/anaconda3/envs/ace_step/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 969, in _run
[rank0]:     self.__setup_profiler()
[rank0]:   File "/home/meow/anaconda3/envs/ace_step/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 1102, in __setup_profiler
[rank0]:     self.profiler.setup(stage=self.state.fn, local_rank=local_rank, log_dir=self.log_dir)
[rank0]:   File "/home/meow/anaconda3/envs/ace_step/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 1264, in log_dir
[rank0]:     dirpath = self.strategy.broadcast(dirpath)
[rank0]:   File "/home/meow/anaconda3/envs/ace_step/lib/python3.10/site-packages/pytorch_lightning/strategies/ddp.py", line 307, in broadcast
[rank0]:     torch.distributed.broadcast_object_list(obj, src, group=_group.WORLD)
[rank0]:   File "/home/meow/anaconda3/envs/ace_step/lib/python3.10/site-packages/torch/distributed/c10d_logger.py", line 81, in wrapper
[rank0]:     return func(*args, **kwargs)
[rank0]:   File "/home/meow/anaconda3/envs/ace_step/lib/python3.10/site-packages/torch/distributed/distributed_c10d.py", line 3483, in broadcast_object_list
[rank0]:     broadcast(object_sizes_tensor, src=global_src, group=group)
[rank0]:   File "/home/meow/anaconda3/envs/ace_step/lib/python3.10/site-packages/torch/distributed/c10d_logger.py", line 81, in wrapper
[rank0]:     return func(*args, **kwargs)
[rank0]:   File "/home/meow/anaconda3/envs/ace_step/lib/python3.10/site-packages/torch/distributed/distributed_c10d.py", line 2714, in broadcast
[rank0]:     work = group.broadcast([tensor], opts)
[rank0]: torch.distributed.DistBackendError: NCCL error in: /pytorch/torch/csrc/distributed/c10d/NCCLUtils.cpp:77, unhandled cuda error (run with NCCL_DEBUG=INFO for details), NCCL version 2.26.2
[rank0]: ncclUnhandledCudaError: Call to CUDA function failed.
[rank0]: Last error:
[rank0]: Cuda failure 2 'out of memory'
[rank0]:[W510 20:10:51.192609090 ProcessGroupNCCL.cpp:1476] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())