redhat-et / ilab-on-ocp

10 stars 14 forks source link

Switch to FSDP #51

Open MichaelClifford opened 1 week ago

MichaelClifford commented 1 week ago

FSDP is toolkit for distributed model training. It is an alternative to Deepspeed. The InstructLab team has added support for FSDP in addition to DeepSpeed in their training repo and we would like to implement FSDP instead of DeepSpeed here.

Please see https://github.com/instructlab/training/pull/235 for the relevant changes.

Hopefully this change should not impact our code base significantly, and should primarily consist of updating our TrainingArgs here: https://github.com/redhat-et/ilab-on-ocp/blob/8748546911c5188ca392a4166d7a2a5231b2b4e3/training/run_main_ds.py#L153-L154

Since this script is included in our PytorchJob image, the image will also need to be rebuilt to fully implement this change.

astefanutti commented 3 days ago

@MichaelClifford Do you think I can help on this one?

MichaelClifford commented 2 days ago

@astefanutti Thanks. I don't think this issue is relevant to the standalone script we are working on since its about modifying a custom run_main_ds.py script that wont be included in (at least the first iteration of) the standalone script. If you would like to help with the current training work, please take a look at #74, we can likely incorporate FSDP implementation in this PR.

astefanutti commented 1 day ago

@MichaelClifford thanks, that's a bit clearer to me now :)

Has the switch to FSDP already be completed in the training repository, or there are still some work to do?

MichaelClifford commented 1 day ago

@astefanutti yes, it has recently been added to the training repo

astefanutti commented 1 day ago

@MichaelClifford awesome, I've just seen #81, you guys are too fast 😃!