Scalable and Efficient Fine-Tuning of LLM on Azure ML

https://github.com/james-tn/llm-fine-tuning/tree/main/opensource_llm/single_step

Faster Iterations, Shorter Time-to-Value:

In today’s competitive AI landscape, time is of the essence. The faster you can fine-tune a model, the quicker you can validate ideas, test hypotheses, and bring solutions to market.

High-profile GPU machines are costly:

High-performance GPUs and compute clusters don’t come cheap, and their availability is often limited. Efficient fine-tuning techniques, such as model sharding and distributed training, maximize the utilization of these precious resources—ensuring that you get the most out of your infrastructure investment.

Not all GPU computes are created equal, and choosing the right sku can make or break your training efficiency.

• ND Series: Ideal for distributed training across multiple nodes, thanks to its Infiniband (IB) connectivity that ensures high-speed communication between nodes like pretraining LLM or finetuning very large model ~70B params.
• NC Series: Small and medium workload where no heavy interaction between nodes needed like LLM inferencing or mid-size LLM finetuning.

Azure GPU Machine Options by Scenario:

When scaling fine-tuning tasks, choosing the right distributed training method is key:

• DDP (Data Parallelism): Works well when the entire model fits on a single GPU. It replicates the model across multiple GPUs and splits the data for parallel processing. Check experiment 1 in the following section.
• Model Parallelism: A game-changer for massive models that don’t fit on a single GPU. It shards not only the data but also the model parameters and optimizer states across multiple GPUs, enabling efficient training of models like LLaMA-70B on GPUs with low memory GPUs. Both FSDP and DeepSpeed as libraries excel at implementing advanced forms of model parallelism and memory optimization.
• Memory Optimization Techniques
  • Gradient Checkpointing: Reduces memory by recomputing activations during the backward pass, trading memory for additional computation.

• • Mixed Precision Training: Reduces memory usage by using FP16 or BF16 instead of FP32, accelerating training while maintaining numerical stability.
    Supported by both frameworks.
  • Quantization (DeepSpeed Exclusive): Uses INT8 precision for weights and activations, dramatically reducing memory and compute requirements.
  • Offloading (DeepSpeed Exclusive): Offloads optimizer states and model parameters to CPU or NVMe, freeing up GPU memory for computation.

We conducted an experiment to fine-tune the Llama-3.1-8B model using LoRA (Low-Rank Adaptation) on Azure ML NDv2-V100 nodes. The goal was to evaluate the efficiency of fine-tuning across different numbers of nodes (1, 2, and 3) and observe the impact on training time and throughput.

Azure ML Job YAML Definition

$schema: https://azuremlschemas.azureedge.net/latest/commandJob.schema.json  
type: command  
code: ./ # Path to your training script and related files  
inputs:
  model_dir: 
    path: azureml://registries/azureml/models/mistralai-Mistral-7B-v01/versions/19
command: >  
  accelerate launch 
  --num_processes 16 # gpu per machine * num of machines 
  --num_machines 2 
  --machine_rank $NODE_RANK
  --main_process_ip $MASTER_ADDR
  --main_process_port $MASTER_PORT
compute: azureml:ndv2-cluster  
resources:  
  instance_count: 2  # Number of nodes for distributed training  
distribution:  
  type: pytorch  
  process_count_per_instance: 1  # Number of processes per node  

Results:

As you increased the number of nodes from one to three, the throughput increased proportionally. This indicates that the system scaled efficiently with the addition of more nodes, maintaining a close-to-linear improvement in throughput.

Fine-tuning a 70B-parameter model on GPUs with only 16GB of memory might sound impossible, but we made it happen using FSDP (Full Sharded Data Parallelism) on Azure ML using a cluster of multiple NDv2-V100 nodes. By distributing not only the data but also the model parameters and optimizer states across multiple nodes, we unlocked the power of full sharding.

 $schema: https://azuremlschemas.azureedge.net/latest/commandJob.schema.json  
type: command    
code: ./ # Path to your training script and related files  
inputs:  
  model_dir:  
    path: azureml://registries/azureml-meta/models/Llama-3.3-70B-Instruct/versions/4  
  
command: >  
  accelerate launch  
  --config_file "configs/fsdp_config.yaml"  
  --num_processes 32 
  --num_machines 4  
  --machine_rank $NODE_RANK  
  --main_process_ip $MASTER_ADDR  
  --main_process_port $MASTER_PORT  
  train.py   

compute: azureml:ndv2-cluster  
resources:  
  instance_count: 4  # Number of nodes for distributed training  
distribution:  
  type: pytorch  
  process_count_per_instance: 1  # Number of processes per node  

Key Takeaways:

• Memory Efficiency: Full sharding enabled us to fine-tune the LLaMA-70B model on V100 GPUs despite their limited memory.
• Connectivity Matters: The Infiniband (IB) connectivity of ND nodes played a critical role in ensuring smooth communication across GPUs, making this feat possible.

Scalable and efficient fine-tuning is the key to unlocking the true potential of Large Language Models. By leveraging distributed training techniques, such as FSDP and DDP, and optimizing compute resources on Azure ML, researchers and practitioners can overcome the challenges of training massive models—reducing costs, accelerating time-to-value, and driving AI innovation.

Access the code and start experimenting here!

Future work:
The second part will focus on real-world pipeline setups, including end-to-end model training, hyperparameter optimization, and testing. The third part will dive into deploying trained models for practical use. Future posts may explore best practices for specific fine-tuning scenarios and techniques.