🤝 Compatibility of the TRL CLI with accelerate arguments

[qgallouedec]

This PR adds two features

• Use any accelerate launch argument in the TRL CLI. For example, now you ca use:

  trl sft \
    --model_name_or_path Qwen/Qwen2.5-0.5B \
    --dataset_name stanfordnlp/imdb \
    --num_processes 4
  

• Use predefined accelerate config, example:

  trl sft \
    --model_name_or_path Qwen/Qwen2.5-0.5B \
    --dataset_name stanfordnlp/imdb \
    --accelerate_config deepspeed_zero2
  

The two above settings stay compatible with config file:

# sft_config.yaml
model_name_or_path: Qwen/Qwen2.5-0.5B
dataset_name: stanfordnlp/imdb
accelerate_config: deepspeed_zero2  # or path/to/my/accelerate/config.yaml

trl sft --config sft_config.yaml

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[lewtun]

Nice QoL improvement! LGTM with some nits

[lewtun]

The QLoRA distinction is not really needed IMO and is an artifact from when we tuned L3 405B. I suggest we merge this #3317 first and then just have fsdp1 and fsdp2

[qgallouedec]

Sounds good, I just rebased on fix-fsdp2, and making the required changes

[qgallouedec]

FSDP1/2 added in 2a52ee6

[shirinyamani]

@qgallouedec LGTM, I only think we are just missing the config support for multi-node, (aka --num_machines as an accelerate arg) which would be sth like;

compute_environment: LOCAL_MACHINE
debug: false
distributed_type: MULTI_GPU
downcast_bf16: 'no'
gpu_ids: all
machine_rank: 0  
main_training_function: main
mixed_precision: 'bf16'
num_machines: 2  
num_processes: 8  
rdzv_backend: static
rdzv_endpoint: "MASTER_NODE_IP:PORT" 
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

but im just not sure abt;

1. defining the IP address of the main node and available port --> rdzv_endpoint
2. also the machine_rank which will be different based of the nodes, i.e. for each machine (0, 1, 2, etc.)

[qgallouedec]

You're right @shirinyamani, we don’t have general multi-node examples because you need to provide things like IP addresses manually. However, I think it should be doable by combining, for example, a Zero3 configuration with these arguments directly in the Slurm script. It would look something like this:

# sft_config.yaml
model_name_or_path: Qwen/Qwen2.5-0.5B
dataset_name: stanfordnlp/imdb
accelerate_config: deepspeed_zero3

#!/bin/bash
#SBATCH --nodes=4
#SBATCH --gres=gpu:8
#SBATCH --ntasks-per-node=8

# Get the list of allocated nodes and set main process IP
NODELIST=($(scontrol show hostnames $SLURM_JOB_NODELIST))
MASTER_ADDR="${NODELIST[0]}"

# Launch distributed training
trl sft sft_config.yaml \
    --num_processes $SLURM_NTASKS \
    --num_machines $SLURM_JOB_NUM_NODES \
    --main_process_ip $MASTER_ADDR \
    --machine_rank $SLURM_PROCID \
    --rdzv_backend c10d

I haven't tested though, but feel free to add this to the the doc in a follow-up PR