Any plans for multi-gpu training support?

[joeljang]

Hey Genmo team, thank you so much for the open-source of the fine-tuning code.

Are there any plans for releasing a multi-gpu fine-tuning version?
Was there a specific reason there is only a single-gpu fine-tuning code release? Is it nontrivial to make it multi-gpu training?

Thanks in advance!

[paras-genmo]

Hi @joeljang, at the moment we're not planning to release a multi-GPU trainer. One of the top community requests has been to reduce the minimum resource requirements for using and fine-tuning Mochi, so we made sure the LoRA tuner could run on a single GPU.

[vedantroy]

@joeljang I don't think it should be terribly difficult to do. You'll need to implement the context parallel backwards passes: https://github.com/genmoai/mochi/blob/main/src/genmo/mochi_preview/dit/joint_model/context_parallel.py + tweak the training loop a bit.

[alexnwang]

@vedantroy Following up on this -- wouldn't it just be a matter of adding DDP? As forward + backward fit fully on a single GPU, shouldn't of sharding data and averaging gradient updates be sufficient?