Switch to FSDP

[MichaelClifford]

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 instructlab/training#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]

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

[MichaelClifford]

@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]

@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]

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

[astefanutti]

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

[MichaelClifford]

closed by #81