feat: move to accelerate launch for distributed training

[kmehant]

Description of the change

Related issue number

Closes #87

How to verify the PR

I have launched a multi gpu training using FSDP and accelerate

accelerate launch --main_process_ip $MASTER_ADDR --main_process_port $MASTER_PORT --config_file config/accelerate_fsdp_config.yaml tuning/sft_trainer.py --peft_method none --tokenizer_name_or_path $MODEL_PATH --model_name_or_path ${MODEL_PATH} --data_path ${DATA_PATH} --output_dir ${OUTPUT_PATH} --num_train_epochs 5 --per_device_train_batch_size 2 --per_device_eval_batch_size 4 --evaluation_strategy "no" --save_strategy "epoch" --learning_rate "1e-5" --warmup_ratio 0.03 --lr_scheduler_type "cosine" --logging_steps 1 --include_tokens_per_second --packing False --response_template "\n### Label:" --dataset_text_field "output" --use_flash_attn False --torch_dtype bfloat16

$MASTER_ADDR and $MASTER_PORT are set by the job launcher (Pytorch Job launcher from kubeflow operator). I have used twitter_complaints.json for the training data and Llama 7B as the base model. Rest of the training arguments used are shown in the above command.

I have used the accelerate config thats being added by this PR (config/accelerate_fsdp_llama_2_procs.yaml).

Some screenshots on the training.

GPU and memory utilization

Training loss curve

Concerns

Saving model

ref - https://huggingface.co/blog/ram-efficient-pytorch-fsdp

set full state dict before saving the model

    trainer.accelerator.state.fsdp_plugin.set_state_dict_type("FULL_STATE_DICT")

Was the PR tested

• I have added >=1 unit test(s) for every new method I have added.
• I have ensured all unit tests pass

[Ssukriti]

thank for the PR!

Minor suggestions , unfortunately we had a duplicate created at same time #91

Fabian had also volunteered earlier to add it, but I thought he was busy and then you both crated PR at same time.

To keep things fair, do you mind co-commiting so you both get credit? You could rebase commits or open a new PR

I reviewed both your PRs and added content here that was different and could be useful.

[kmehant]

@Ssukriti I have addressed your comments and cherry picked Fabian's commit from his PR to here (automatically credits him at commit 4aa0153). Looking forward to your review, thanks.

[fabianlim]

@kmehant @Ssukriti I have some suggestions. I feel the PR can be further streamlined.

[Ssukriti]

will approve and merge when Fabian's comments are addressed. I had another small comment on the name of training_data_path variable.

Just on the environment variables , I would suggest keeping them as examples to avoid confusion on how to set them.

`# Please set the environment variables:
# MODEL_PATH = llama-7b-hf # Huggingface model id or path to a checkpoint
# TRAIN_DATA_PATH=twitter_complaints.json # Path to the train dataset
# OUTPUT_PATH=out # Path to the output folder where the checkpoints are saved`

Will wait for other comments from Fabian to be addressed. Thank you!

[kmehant]

@Ssukriti @fabianlim I have collated all your review comments and have updated the PR, please review, thank you.

[kmehant]

@Ssukriti The PR is ready, thank you.

[Ssukriti]

thank you!