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:

ilab-on-ocp/training/run_main_ds.py

Lines 153 to 154 in 8748546

	training_args = TrainingArgs(
	deepspeed_options = DeepSpeedOptions(cpu_offload_optimizer=args.cpu_offload),

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