A Comprehensive Guide to Multi-GPU LoRA Fine-Tuning with Distributed Data Parallelism and Sequence Parallelism using Axolotl — 1

1. Introduction

Large Language Models (LLMs) pre-trained on vast datasets possess remarkable general language understanding capabilities. However, adapting these models to specific downstream tasks or domains often requires further training, a process known as fine-tuning. Full fine-tuning, which involves updating all model parameters, becomes computationally prohibitive for models with billions of parameters, demanding significant memory and compute resources. Parameter-Efficient Fine-Tuning (PEFT) methods, such as Low-Rank Adaptation (LoRA), have emerged as effective alternatives, drastically reducing the number of trainable parameters while maintaining comparable performance.

Training even with PEFT methods can still be time-consuming. Distributed training techniques are employed to accelerate this process by leveraging multiple GPUs. Common strategies include Data Parallelism, Tensor Parallelism, and Pipeline Parallelism. Furthermore, processing very long input sequences, a growing requirement for many applications, introduces significant memory challenges, particularly within the attention mechanism. Sequence Parallelism (SP) offers a solution by partitioning the sequence dimension across multiple devices.

This blog provides a comprehensive technical guide and runbook for fine-tuning the Qwen-2.5–14B-Instruct model using LoRA on a multi-GPU system (specifically, 8 NVIDIA A100–40GB GPUs on a single node). It details a specific strategy combining Distributed Data Parallelism (DDP) with Sequence Parallelism, orchestrated using the Axolotl framework.

2. Core Concepts

2.1 LoRA (Low-Rank Adaptation)

LoRA is a highly effective Parameter-Efficient Fine-Tuning (PEFT) technique designed to adapt large pre-trained models with minimal computational overhead. Instead of updating all the weights of the base model (W₀), LoRA freezes them and injects smaller, trainable “adapter” matrices into specific layers, typically the linear layers within the Transformer architecture. For a given weight matrix W₀ of dimension d×k, LoRA introduces two low-rank matrices, A (dimension r×k) and B (dimension d×r), where the rank r is significantly smaller than d and k (r≪min(d,k)). During training, only matrices A and B are updated. The forward pass is modified such that the output h becomes h=W0​x+ΔWx=W0​x+BAx. The matrix B is initialized with zeros, ensuring the initial state matches the pre-trained model. This approach drastically reduces the number of trainable parameters (by up to 10,000 times for models like GPT-3 175B) and GPU memory requirements (by up to 3 times) compared to full fine-tuning. After training, the learned update ΔW=BA can be merged back into the original weights (W=W0​+BA), resulting in a standard model architecture with no added inference latency, a key advantage over methods like adapter modules. LoRA enables efficient task-switching by simply swapping the small A and B matrices for different tasks.

2.2 DDP (Distributed Data Parallel)

Distributed Data Parallel (DDP) is PyTorch’s recommended method for data parallelism across multiple GPUs, potentially spanning multiple machines. In DDP, the model is replicated on each participating GPU (process). Each GPU process receives a distinct subset (shard) of the input data batch. During the forward pass, each replica processes its data shard independently. In the backward pass, as gradients are computed locally on each GPU, DDP uses torch.distributed communication collectives (primarily all_reduce) to synchronize these gradients across all replicas. This synchronization typically involves averaging the gradients, ensuring that all model replicas perform the same weight update during the optimizer step. DDP achieves this synchronization efficiently by registering autograd hooks that trigger communication as soon as gradients are ready, often overlapping communication with the ongoing backward computation, and using gradient bucketing to further optimize communication bandwidth. Compared to the older torch.nn.DataParallel (DP), DDP is generally faster, especially in multi-GPU single-node setups, due to its multi-process architecture which avoids Python's Global Interpreter Lock (GIL) contention.

2.3 Sequence Parallelism

Sequence Parallelism (SP) is a model parallelism technique specifically designed to address the memory bottleneck associated with processing long input sequences in Transformer models. The primary memory challenge arises from the self-attention mechanism, where the Key-Value (KV) cache and intermediate activation sizes scale quadratically or linearly with sequence length.1 SP mitigates this by partitioning the input sequence along its length dimension and distributing these segments across multiple GPUs. Each GPU then processes only its assigned portion of the sequence for certain operations. Techniques like Ring Attention implement SP by combining blockwise computation of attention with efficient ring-based communication protocols. In Ring Attention, devices arranged in a conceptual ring exchange necessary KV blocks (representing other parts of the sequence) while computing attention on their local block, crucially overlapping communication with computation to minimize latency overhead. This allows the effective context length to scale almost linearly with the number of devices, enabling the processing of sequences far exceeding the memory capacity of a single GPU. Implementations often leverage optimized kernels like FlashAttention to further enhance memory efficiency and speed within each block’s computation.

2.4 Axolotl

Axolotl is an open-source tool designed to streamline and simplify the fine-tuning process for various LLMs. It acts as a high-level wrapper around underlying libraries like Hugging Face Transformers, PEFT, PyTorch, DeepSpeed, and Accelerate. Axolotl’s core strength lies in its use of YAML configuration files, which allow users to define the entire fine-tuning workflow — including model selection, dataset specification, training method (full fine-tuning, LoRA, QLoRA), optimization parameters, distributed training setup (DDP, FSDP, DeepSpeed, SP), and performance enhancements (Flash Attention, packing) — without writing extensive boilerplate code. It supports a wide range of models and dataset formats, integrates best practices, and facilitates common tasks like preprocessing, training, inference, and merging LoRA adapters. This configuration-driven approach makes it easier to experiment with different setups and reproduce results.

3. Rationale for DDP + Sequence Parallelism Strategy

The choice of combining Distributed Data Parallelism (DDP) with Sequence Parallelism (SP) for LoRA fine-tuning Qwen-2.5–14B on a multi-A100 setup is driven by several factors related to efficiency and the specific characteristics of LoRA:

1. LoRA Reduces Optimizer State: LoRA significantly reduces the number of trainable parameters compared to full fine-tuning. Consequently, the memory required for optimizer states (like momentum and variance in AdamW) and gradients is substantially smaller. Strategies like Fully Sharded Data Parallelism (FSDP) or DeepSpeed ZeRO (especially Stage 3) primarily offer memory savings by sharding these optimizer states, gradients, and potentially the model parameters themselves across GPUs. Since LoRA already makes these components small, the primary benefits of FSDP/ZeRO Stage 3 become less critical. While ZeRO Stage 1 or 2 could still offer some optimizer state sharding benefits, standard DDP often provides sufficient data throughput for LoRA training.
2. Sequence Length is the Bottleneck: When fine-tuning with long sequences (e.g., 4096 tokens in this case), the dominant memory consumer is often the activation memory, particularly the KV cache within the attention mechanism, which scales with sequence length. DDP alone does not alleviate this per-GPU activation memory pressure, as each GPU holds the full model replica and processes sequences of the full length. Sequence Parallelism, however, directly addresses this by partitioning the sequence dimension itself across GPUs. This partitioning reduces the sequence length handled by each GPU within the attention computation, thereby lowering the activation memory footprint per device.
3. SP Complements DDP: DDP provides parallelism across the batch dimension, increasing data throughput, while SP provides parallelism across the sequence dimension, enabling longer contexts. They operate orthogonally and can be combined. In this setup, DDP handles the distribution of batches across the 8 GPUs, and SP (with degree 4) further splits the attention computation for each sequence within a batch across rings of 4 GPUs (GPUs 0–3 form one ring, GPUs 4–7 form another)
4. Efficiency of Ring Attention: Modern SP implementations like Ring Attention (leveraged by Axolotl via ring-flash-attn) are designed to overlap the communication required for SP with the computation of the attention blocks, minimizing the added latency overhead. This makes SP computationally efficient, especially when combined with optimized kernels like FlashAttention.
5. Simplicity and Availability in Axolotl: Axolotl provides straightforward configuration options for enabling both DDP (via Accelerate) and SP (via ring-flash-attn) through its YAML interface. Axolotl version 0.8.1 specifically includes the necessary kernels for SP with Qwen-2.5. This avoids the need for complex manual implementation of these parallelism strategies.

4. Environment Setup Steps

The runbook targets a high-performance computing node with the following specifications:

GPUs: 8 x NVIDIA A100-SXM4–40GB GPUs interconnected within a single machine. The A100, based on the Ampere architecture (SM80), provides strong performance for mixed-precision computation (BF16/FP16) and features high-bandwidth NVLink interconnects crucial for efficient multi-GPU communication (like DDP’s all-reduce and SP’s ring communication).

The following commands, executed from a fresh login prompt, establish the necessary environment:

1. Install OS Packages: Update package lists and install essential build tools, version control for large files (git-lfs), and utilities like tmux.

sudo apt-get update -y
sudo apt-get install -y git-lfs tmux build-essential ninja-build cmake

These packages provide compilers (build-essential), build systems (cmake, ninja-build), and version control (git-lfs) needed for installing Python packages with C++/CUDA extensions from source, which is required for libraries like ring-flash-attn.

2. Create Python Virtual Environment: Isolate project dependencies using venv

python3 -m venv axo-env
source axo-env/bin/activate
pip install --upgrade pip setuptools wheel

3. Install PyTorch: Install PyTorch version 2.5.1 compiled for CUDA 12.1 using the official PyTorch wheel index.

pip install --index-url https://download.pytorch.org/whl/cu121 \
            torch==2.5.1+cu121

PyTorch 2.5 is chosen for its stable bfloat16 support (beneficial on A100s) and potential improvements or fixes in underlying communication libraries like NCCL compared to earlier versions.46

4. Install Axolotl with Sequence Parallel Kernels: Install Axolotl version 0.8.1, specifically including the ring-flash-attn extras which provide the sequence parallelism capabilities. Set the TORCH_CUDA_ARCH_LIST environment variable before installation.

export TORCH_CUDA_ARCH_LIST="8.0"           # A100’s SM80
pip install --no-build-isolation \
    "axolotl[ring-flash-attn]==0.8.1"

• TORCH_CUDA_ARCH_LIST="8.0": This environment variable instructs the PyTorch build process (specifically for extensions compiled during installation, like the SP kernels) to generate code specifically for the CUDA Compute Capability 8.0 (Ampere architecture, used by A100 GPUs). Pre-compiling for the target architecture avoids Just-In-Time (JIT) compilation during the first run, preventing potential startup delays or stalls.
• axolotl[ring-flash-attn]==0.8.1: Axolotl version 0.8.1 is the first tagged release bundling the necessary ring-flash-attn v2 kernels required for sequence parallelism with models like Qwen-2.5 on 4-8 GPUs. The [ring-flash-attn] extra ensures these specific dependencies are installed.
• --no-build-isolation: This flag tells pip to build the package in the current environment rather than an isolated one, which can sometimes be necessary when dependencies have complex build requirements or rely on environment variables set beforehand (like TORCH_CUDA_ARCH_LIST).

This setup ensures that all required libraries are installed with compatible versions and compiled appropriately for the A100 hardware, laying the foundation for the distributed fine-tuning process.

5. Dataset Preparation and Configuration

5.1 Dataset Selection and Transformation

The provided runbook utilizes the ccdv/govreport-summarization dataset, a publicly available corpus containing government reports and their corresponding summaries. For instruction fine-tuning, the goal is to teach the model to follow a specific instruction (e.g., "Summarize...") given some input context. The chosen format is "Alpaca-style" JSON Lines (JSONL), which requires specific keys for instruction, input, and output.

The following Python script, executed once, performs this transformation:

from datasets import load_dataset
import pathlib, json, yaml

# Load the dataset (only train split used here)
ds = load_dataset("ccdv/govreport-summarization", split="train")

# Define a function to map original columns to Alpaca format
def _map(x):
    return {
        "instruction": "Summarize the following report.", # Fixed instruction
        "input": x["report"], # Use report as input for the instruction
        "output": x["summary"], # Use summary as the desired output
    }

# Apply the mapping and remove original columns
ds = ds.map(_map, remove_columns=ds.column_names)

# Write the transformed data to a JSONL file
pathlib.Path("train_data.jsonl").write_text(
    "\n".join(json.dumps(r) for r in ds), encoding="utf-8"
)

This script loads the dataset using the datasets library, defines a mapping function _map that restructures each example into a dictionary with keys instruction, input, and output, applies this transformation to the entire dataset, and saves the result as train_data.jsonl. Each line in this file is a JSON object representing one training example.

5.2 Importance of Data Formatting and Quality

The structure of the data in train_data.jsonl is critical. Axolotl uses the specified type (e.g., alpaca) in its configuration to understand how to parse these JSON objects and construct the final input sequence for the model during training. It typically concatenates the instruction, input, and output fields (potentially using a predefined template) into a single string. Crucially, Axolotl also uses this structure to determine which parts of the generated sequence should contribute to the loss calculation. For instruction tuning, the loss is usually calculated only on the tokens corresponding to the output field, teaching the model to generate the desired response given the instruction and input. Mismatches between the actual data format and the type specified in the config, or errors in the data itself, can lead to incorrect training signals and poor model performance.

5.3 Axolotl Configuration Specification

The prepared dataset is specified within the Axolotl YAML configuration file (qwen14b_lora_4096_ddp.yaml) under the datasets key:

datasets:
  - path: train_data.jsonl # Local path to the prepared JSONL file
    ds_type: json          # Specifies the file format is JSON (implicitly JSONL)
    type: alpaca             # Specifies the instruction format structure

• path: Points to the location of the dataset file. Axolotl supports local paths, Hugging Face dataset identifiers (e.g., teknium/GPT4-LLM-Cleaned), and potentially other sources.34
• ds_type: Indicates the data format. json is used for JSON Lines files. Other types might exist depending on the Axolotl version.
• type: Specifies the structure or template Axolotl should expect within each JSON object. alpaca corresponds to the {"instruction":..., "input":..., "output":...} structure.39 Axolotl supports various other types like sharegpt, oasst, completion, etc., each expecting different fields or formats.39

This configuration tells Axolotl how to find, load, and interpret the training data according to the Alpaca instruction-following format.

6. Axolotl Configuration Deep Dive

The Axolotl YAML configuration file serves as the central control panel for the entire fine-tuning process, encapsulating all hyperparameters and settings.

# qwen14b_lora_4096_ddp.yaml
base_model: Qwen/Qwen2.5-14B-Instruct
model_type: AutoModelForCausalLM
tokenizer_type: AutoTokenizer

adapter: lora
lora_r: 16
lora_alpha: 16
lora_dropout: 0.05
lora_target_linear: true

bf16: true
load_in_8bit: true
gradient_checkpointing: true

sequence_len: 4096
flash_attention: true
sequence_parallel_degree: 4     #if error, replace it with 1 # 1 ← OFF
# ring_attn_func: batch_ring      # ← REMOVE

datasets:
  - path: /home/ubuntu/train_data.jsonl
    ds_type: json
    type: alpaca

output_dir: ./outputs/qwen14b-4096-lora-ddp
micro_batch_size: 2              # 
gradient_accumulation_steps: 8  # if results in error, replace it with 4
num_epochs: 2
learning_rate: 2.0e-4
lr_scheduler: cosine
warmup_steps: 100
optimizer: adamw_bnb_8bit

6.1 Base Model & Tokenizer

These parameters define the starting point for the fine-tuning process.

• base_model: Qwen/Qwen2.5-14B-Instruct: Specifies the Hugging Face Hub identifier for the pre-trained model. Qwen-2.5-14B-Instruct is a 14-billion parameter model already tuned for following instructions. Choosing an appropriate base model aligned with the target task's nature (e.g., instruction-following, chat, base language modeling) is crucial for successful fine-tuning.2
• model_type: AutoModelForCausalLM: Instructs the underlying Hugging Face Transformers library to automatically load the correct Python class (Qwen2ForCausalLM in this case) designed for causal language modeling tasks (i.e., predicting the next token).
• tokenizer_type: AutoTokenizer: Similarly, directs Transformers to load the tokenizer specifically associated with the base_model. Using the correct tokenizer is essential as it defines how text is converted into the token IDs the model understands.4 Mismatches can lead to nonsensical inputs and outputs.

6.2 Data & Output

These parameters specify the data source and where results should be stored.

• datasets: Configures the training dataset(s). As detailed in Section 5.3, this includes the path, file type, and instruction format.

datasets:
  - path: train_data.jsonl
    ds_type: json
    type: alpaca

• output_dir:./outputs/qwen14b-4096-lora-ddp: Defines the directory where Axolotl will save training outputs, including model checkpoints (containing the LoRA adapter weights), training logs, and potentially the final merged model if requested.

6.3 Optimization

These settings control the learning process itself.

• optimizer: adamw_bnb_8bit: Selects the AdamW optimizer implemented with 8-bit states provided by the BitsAndBytes library. AdamW incorporates weight decay correction compared to standard Adam. The 8-bit quantization significantly reduces the memory required to store optimizer states (e.g., momentum, variance estimates), often halving it compared to 16-bit or 32-bit states, with typically negligible impact on final model quality. This is a common choice for memory-constrained scenarios, especially when using LoRA or QLoRA.
• learning_rate: 2.0e-4: Sets the initial learning rate (step size) for the optimizer. The optimal learning rate is task- and model-dependent and often requires tuning.2 A rate of 2.0×10−4 is within the typical range explored for LoRA fine-tuning.
• lr_scheduler: cosine: Specifies the learning rate schedule. A cosine scheduler starts at the initial learning_rate, gradually decreases it following a cosine curve towards zero over the course of training. This often helps the model converge more smoothly and potentially reach a better minimum compared to a constant learning rate.
• warmup_steps: 100: Implements a learning rate warmup phase. For the first 100 training steps, the learning rate linearly increases from 0 to the target learning_rate (2.0e-4). This helps stabilize training in the initial stages when parameters are changing rapidly.
• num_epochs: 2: Defines the total number of times the training process will iterate over the entire dataset . For PEFT methods like LoRA, fewer epochs (e.g., 1-3) are often sufficient to adapt the model compared to full fine-tuning.

6.4 Batching

Batching parameters determine how data is fed to the GPUs and when model weights are updated.

• micro_batch_size: 1: This is the number of data samples processed by each individual GPU in a single forward and backward pass. A value of 1 means each of the 8 GPUs processes one sequence at a time before gradient synchronization (for DDP) or other parallel computations (for SP). This value is often constrained by GPU memory, especially with long sequences.
• gradient_accumulation_steps: 8: Instructs the training loop to perform 8 forward and backward passes (using micro-batches) and accumulate the gradients before executing a single optimizer step (weight update). This effectively simulates a larger batch size without requiring the corresponding memory.
• Global Batch Size: The effective batch size used for each weight update is calculated as: GlobalBatchSize=micro_batch_size×num_gpus×gradient_accumulation_steps
• In this configuration: 1×8×8=64. This means the model weights are updated based on the averaged gradients from 64 sequences, even though only 1 sequence fits onto each GPU at a time per pass. Maintaining an adequate global batch size is important for stable training convergence.
• gradient_checkpointing: True: Enables gradient checkpointing (also known as activation checkpointing). Instead of storing all intermediate activations from the forward pass (which consumes significant memory for deep models and long sequences), gradient checkpointing discards them and recomputes them during the backward pass when needed for gradient calculation. This significantly reduces activation memory usage (estimated ~35% VRAM saving in this setup) at the cost of increased computational time due to the recomputation. It's often essential for training with long sequences or large models within limited VRAM.

6.5 Context & Attention

These parameters are critical for handling long sequences and optimizing the attention mechanism.

• sequence_len: 4096: Sets the maximum sequence length (in tokens) that the model will process during training. This directly impacts the memory required for activations and the KV cache. The value of 4096 is chosen as a balance that fits within the 40GB A100 memory budget given the other memory-saving techniques employed (8-bit loading, bf16, LoRA, gradient checkpointing, SP).
• flash_attention: True: Enables the use of FlashAttention (likely FlashAttention-2 given the context and PyTorch version). FlashAttention is a highly optimized attention implementation that computes attention using tiling and reduces memory reads/writes between GPU high-bandwidth memory (HBM) and on-chip SRAM.31 It significantly reduces memory usage (linear instead of quadratic in sequence length for intermediate matrices) and speeds up computation, making it almost mandatory for long sequence lengths. It's also a prerequisite for the sequence parallelism implementation used here.
• sequence_parallel_degree: 4: Sets the degree for sequence parallelism. This means each input sequence will be conceptually split into 4 chunks, and the attention computation for that sequence will be distributed across a ring of 4 GPUs. With 8 GPUs total, Axolotl implicitly creates two independent SP rings (e.g., GPUs 0-3 and GPUs 4-7). This reduces the memory burden related to sequence length on each GPU by approximately a factor of 4 for the attention mechanism, enabling longer sequences. The degree must be a divisor of the total number of GPUs (8 in this case) and ideally also compatible with the number of attention heads and sequence length for optimal load balancing, although the ring-flash-attn library aims to handle this.
• ring_attn_func: "batch_ring": Specifies the specific kernel function from the ring-flash-attn library to use for the sequence-parallel attention computation. batch_ring is one of the available implementations, suitable for standard batched inputs where sequences are not packed. Other options like zigzag_ring_flash_attn_varlen_func or llama3_flash_attn_varlen_func might offer different trade-offs or be better suited for packed sequences.

6.6 Precision & Quantization

These settings control the numerical precision used during training and for model weights, impacting memory usage and computational speed.

• bf16: True: Enables automatic mixed-precision training using the bfloat16 format for activations and computations where possible. Bfloat16 (Brain Floating Point) uses 16 bits like fp16 but allocates more bits to the exponent, providing a wider dynamic range closer to fp32.This makes it less prone to overflow/underflow issues compared to fp16, especially during long training runs, while still offering significant memory savings (activations stored in 16 bits instead of 32) and speedups (leveraging Tensor Cores on A100s). PyTorch 2.5 provides stable support for bf16.
• load_in_8bit: True: Instructs Axolotl (via the Transformers and BitsAndBytes libraries) to load the weights of the base_model quantized to 8-bit integers. This reduces the memory required to store the frozen base model parameters by approximately 4x compared to fp32 or 2x compared to fp16/bf16.54 The computations involving these weights are typically performed in higher precision (like bf16 here) after de-quantization within the compute kernels. This is crucial for fitting large base models like Qwen-14B into 40GB VRAM alongside activations for long sequences.

6.7 LoRA Adapter

These parameters configure the LoRA-specific aspects of the fine-tuning.

• adapter: lora: Explicitly selects LoRA as the PEFT method. Axolotl also supports other adapters like qlora.
• lora_r: 16: Sets the rank r of the LoRA decomposition matrices (A and B). This is a critical hyperparameter controlling the capacity (number of trainable parameters) of the adapter. A higher rank means more parameters but potentially better adaptation capability, while a lower rank is more parameter-efficient. Rank 16 is presented as an empirical sweet spot for this model and task, resulting in approximately 0.46% trainable parameters relative to the base model .
• lora_alpha: 16: Sets the scaling factor α for the LoRA update ΔW=BA.6 The update is often scaled by α/r. A common heuristic is to set α=r, but it can be treated as another hyperparameter to tune.
• lora_dropout: 0.05: Applies dropout with a probability of 0.05 to the LoRA adapter weights (specifically, often applied after the first matrix A but before the second matrix B) during training. This acts as a regularization technique to prevent the adapter from overfitting to the fine-tuning data.
• lora_target_linear: True: A convenience flag in Axolotl that automatically targets all modules of type torch.nn.Linear within the base model for LoRA adaptation. This typically includes the query, key, value, and output projection layers in the attention blocks, as well as the layers in the feed-forward networks. Alternatively, one could specify a list of module names using lora_target_modules for more granular control.

7. Distributed Training Setup with Hugging Face Accelerate

To execute the training across the 8 A100 GPUs using the DDP strategy, the Hugging Face Accelerate library is employed. Axolotl leverages Accelerate for launching and managing distributed training jobs.36

7.1 Role of Accelerate

Accelerate provides a unified interface to simplify the process of writing and launching PyTorch code on various distributed hardware configurations, including multi-GPU setups on a single node.58 It abstracts away much of the boilerplate code associated with initializing distributed process groups, handling device placement, and wrapping models, optimizers, and dataloaders for distributed training frameworks like DDP, FSDP, or DeepSpeed.58 By using Accelerate, the same training script can often be run across different setups with minimal or no changes.60

7.2 Interactive Configuration (accelerate config)

Before launching a distributed job, Accelerate needs information about the training environment. The accelerate config command provides an interactive way to specify this setup. For the runbook's scenario (8 GPUs on a single machine using DDP with bf16), the interaction would typically proceed as follows:

$ accelerate config
------------------------------------------------------------------------------------------------------------------------------------------
In which compute environment are you running?
This machine                                                                                                                              
------------------------------------------------------------------------------------------------------------------------------------------
Which type of machine are you using?
multi-GPU                                                                                                                                 
------------------------------------------------------------------------------------------------------------------------------------------
How many different machines will you use (use more than 1 for multi-node training)? [1]: 1                                                
------------------------------------------------------------------------------------------------------------------------------------------
How many GPU(s) should be used for distributed training? [Enter all GPU indices to use separated by commas or 'all'][all]: 8             
------------------------------------------------------------------------------------------------------------------------------------------
Do you wish to use DeepSpeed? [yes/NO]: NO                                                                                                
------------------------------------------------------------------------------------------------------------------------------------------
Do you wish to use Fully Sharded Data Parallel (FSDP)? [yes/NO]: NO                                                                       
------------------------------------------------------------------------------------------------------------------------------------------
Do you wish to use Megatron-LM? [yes/NO]: NO                                                                                             
------------------------------------------------------------------------------------------------------------------------------------------
Do you want to use automatic mixed precision training? [no/fp16/bf16/fp8]: bf16                                                            
------------------------------------------------------------------------------------------------------------------------------------------
Do you want to use TorchDynamo? [yes/NO]: NO                                                                                              
------------------------------------------------------------------------------------------------------------------------------------------

(Note: Specific prompts might vary slightly between Accelerate versions. The prompt regarding TorchDynamo might appear depending on the version.)

This interactive session generates a configuration file, typically saved as default_config.yaml in the Accelerate cache directory (e.g., ~/.cache/huggingface/accelerate/). This file stores the user's choices.

7.3 Significance of Configuration Settings

When accelerate launch is executed, it reads this configuration and automatically handles the initialization of the PyTorch distributed environment (torch.distributed.init_process_group). For a MULTI_GPU setup on NVIDIA hardware, it typically uses the nccl backend, which is highly optimized for inter-GPU communication over NVLink. It ensures that environment variables like RANK, WORLD_SIZE, and LOCAL_RANK are correctly set for each of the 8 launched processes, enabling them to coordinate as part of the DDP group. This abstraction layer removes the need for manual setup using torchrun or complex shell scripts, but understanding the underlying distributed concepts remains beneficial for debugging potential issues like NCCL errors or hangs.

7.4 Environment Variables for Execution

While accelerate config sets up the distributed launch parameters, specific runtime environment variables might still be needed for optimizing execution. The following is set before the accelerate launch command:

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

This variable tunes PyTorch’s CUDA memory allocator. Setting expandable_segments:True allows the allocator to potentially reduce memory fragmentation by growing allocated memory segments rather than strictly adhering to pre-allocated block sizes. This can be beneficial in scenarios with varying tensor sizes or when memory usage is close to the limit, potentially preventing out-of-memory errors caused by fragmentation rather than true exhaustion.

8. Executing and Monitoring the Fine-tuning Process

With the environment, dataset, and configurations prepared, the fine-tuning process can be initiated and monitored.

8.1 Launching the Training Job

The training job is launched using the Accelerate CLI, which reads the configuration set by accelerate config (or uses defaults/command-line overrides) and executes the specified Python module in a distributed manner.

accelerate launch -m axolotl.cli.train qwen14b_lora_4096_ddp.yaml

• accelerate launch: The command to start the distributed execution environment managed by Accelerate.58 It spawns 8 processes (as configured) across the 8 GPUs.
• -m axolotl.cli.train: Tells Python to run the train submodule within the axolotl.cli package as the main script.36
• qwen14b_lora_4096_ddp.yaml: The argument passed to the Axolotl training script, specifying the configuration file to use for this run.

This command initiates the fine-tuning process, loading the model and data, and beginning the training loop across all 8 GPUs using DDP for data parallelism and the configured SP for attention computation.

8.2 Monitoring Progress

During training, Axolotl logs progress information to the console (you may leverage experiment tracking services like WandB if configured). Key metrics to monitor include:

• Training Logs: The console output typically shows metrics for each logging step (determined by training arguments not explicitly shown in the minimal YAML, but configurable in Axolotl). An example log entry provided is:
• {'loss': 0.97, 'grad_norm': 0.06, 'epoch': 1.94}
• loss: The value of the loss function (e.g., cross-entropy loss for language modeling). It should generally decrease over training, indicating learning. Volatility is expected, but a consistently increasing or stagnant loss might indicate issues with learning rate, data, or configuration.
• grad_norm: The norm of the gradients before or after clipping (if enabled). Monitoring this helps detect exploding gradients. A value of 0.06 suggests gradients are well-behaved in this snapshot.
• epoch: The current epoch number, indicating progress through the dataset.
• Other metrics like learning_rate and steps_per_second are also commonly logged.

9. Post-Training Workflow: Inference, Merging, and Sharing

Once the Axolotl training process completes successfully, several steps can be taken to evaluate, utilize, and share the resulting fine-tuned model adapter.

9.1 Local Inference with LoRA Adapter

Before deploying or sharing, it’s essential to test the fine-tuned model’s performance locally. Since LoRA training only produces adapter weights, inference requires loading the original base model and applying the learned adapter dynamically using the peft (Parameter-Efficient Fine-Tuning) library from Hugging Face.

The following Python script demonstrates this process:

Python

from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel, PeftConfig

BASE = "Qwen/Qwen2.5-14B-Instruct"
ADAPTER = "./outputs/qwen14b-4096-lora-ddp"   # your folder above

# 1. Load tokenizer (local copy is fine, but hub works too)
tok = AutoTokenizer.from_pretrained(ADAPTER, trust_remote_code=True)

# 2. Load base model in 8-bit bfloat16 to save VRAM
model = AutoModelForCausalLM.from_pretrained(
    BASE,
    load_in_8bit=True,
    torch_dtype="auto",
    device_map="auto",        # will pick GPU-0
    trust_remote_code=True,
)

# 3. Merge in your LoRA weights
model = PeftModel.from_pretrained(model, ADAPTER)
model.eval()

# 4. Chat-style prompt (matches your Alpaca template)
instruction = "Summarise the Occupational Safety and Health Act in one paragraph."
# input_txt  = ""             # Alpaca ‘input’ field — empty here
input_txt = """


Public Law 91-596 84 STAT. 1590 91st Congress, S.2193 December 29, 1970, as amended through January 1, 2004. (1) An Act

 

To assure safe and healthful working conditions for working men and women; by authorizing enforcement of the standards developed under the Act; by assisting and encouraging the States in their efforts to assure safe and healthful working conditions; by providing for research, information, education, and training in the field of occupational safety and health; and for other purposes. Be it enacted by the Senate and House of Representatives of the United States of America in Congress assembled, That this Act may be cited as the "Occupational Safety and Health Act of 1970."

Footnote (1) See Historical notes at the end of this document for changes and amendments affecting the OSH Act since its passage in 1970 through January 1, 2004.
SEC.
2.
Congressional Findings and Purpose
(a)

The Congress finds that personal injuries and illnesses arising out of work situations impose a substantial burden upon, and are a hindrance to, interstate commerce in terms of lost production, wage loss, medical expenses, and disability compensation payments.
(b)

The Congress declares it to be its purpose and policy, through the exercise of its powers to regulate commerce among the several States and with foreign nations and to provide for the general welfare, to assure so far as possible every working man and woman in the Nation safe and healthful working conditions and to preserve our human resources --  
(1)

by encouraging employers and employees in their efforts to reduce the number of occupational safety and health hazards at their places of employment, and to stimulate employers and employees to institute new and to perfect existing programs for providing safe and healthful working conditions;
(2)

by providing that employers and employees have separate but dependent responsibilities and rights with respect to achieving safe and healthful working conditions;
(3)

by authorizing the Secretary of Labor to set mandatory occupational safety and health standards applicable to businesses affecting interstate commerce, and by creating an Occupational Safety and Health Review Commission for carrying out adjudicatory functions under the Act;
(4)

by building upon advances already made through employer and employee initiative for providing safe and healthful working conditions;
(5)

by providing for research in the field of occupational safety and health, including the psychological factors involved, and by developing innovative methods, techniques, and approaches for dealing with occupational safety and health problems;
(6)

by exploring ways to discover latent diseases, establishing causal connections between diseases and work in environmental conditions, and conducting other research relating to health problems, in recognition of the fact that occupational health standards present problems often different from those involved in occupational safety;
(7)

by providing medical criteria which will assure insofar as practicable that no employee will suffer diminished health, functional capacity, or life expectancy as a result of his work experience;
(8)

by providing for training programs to increase the number and competence of personnel engaged in the field of occupational safety and health; affecting the OSH Act since its passage in 1970 through January 1, 2004.
(9)

by providing for the development and promulgation of occupational safety and health standards;
(10)

by providing an effective enforcement program which shall include a prohibition against giving advance notice of any inspection and sanctions for any individual violating this prohibition;
(11)

by encouraging the States to assume the fullest responsibility for the administration and enforcement of their occupational safety and health laws by providing grants to the States to assist in identifying their needs and responsibilities in the area of occupational safety and health, to develop plans in accordance with the provisions of this Act, to improve the administration and enforcement of State occupational safety and health laws, and to conduct experimental and demonstration projects in connection therewith;
(12)

 by providing for appropriate reporting procedures with respect to occupational safety and health which procedures will help achieve the objectives of this Act and accurately describe the nature of the occupational safety and health problem;
(13)

 by encouraging joint labor-management efforts to reduce injuries and disease arising out of employment.
SEC.
3.
Definitions

For the purposes of this Act --
(1)

The term "Secretary" means the Secretary of Labor.
(2)

The term "Commission" means the Occupational Safety and Health Review Commission established under this Act.
(3)

The term "commerce" means trade, traffic, commerce, transportation, or communication among the several States, or between a State and any place outside thereof, or within the District of Columbia, or a possession of the United States (other than the Trust Territory of the Pacific Islands), or between points in the same State but through a point outside thereof.
(4)

The term "person" means one or more individuals, partnerships, associations, corporations, business trusts, legal representatives, or any organized group of persons.
(5)

The term "employer" means a person engaged in a business affecting commerce who has employees, but does not include the United States (not including the United States Postal Service) or any State or political subdivision of a State.
(6)

The term "employee" means an employee of an employer who is employed in a business of his employer which affects commerce.
(7)

The term "State" includes a State of the United States, the District of Columbia, Puerto Rico, the Virgin Islands, American Samoa, Guam, and the Trust Territory of the Pacific Islands.
(8)

The term "occupational safety and health standard" means a standard which requires conditions, or the adoption or use of one or more practices, means, methods, operations, or processes, reasonably necessary or appropriate to provide safe or healthful employment and places of employment.
(9)

The term "national consensus standard" means any occupational safety and health standard or modification thereof which (1), has been adopted and promulgated by a nationally recognized standards-producing organization under procedures whereby it can be determined by the Secretary that persons interested and affected by the scope or provisions of the standard have reached substantial agreement on its adoption, (2) was formulated in a manner which afforded an opportunity for diverse views to be considered and (3) has been designated as such a standard by the Secretary, after consultation with other appropriate Federal agencies.
(10)

The term "established Federal standard" means any operative occupational safety and health standard established by any agency of the United States and presently in effect, or contained in any Act of Congress in force on the date of enactment of this Act.
(11)

The term "Committee" means the National Advisory Committee on Occupational Safety and Health established under this Act.
(12)

The term "Director" means the Director of the National Institute for Occupational Safety and Health.
(13)

The term "Institute" means the National Institute for Occupational Safety and Health established under this Act.
(14)

The term "Workmen's Compensation Commission" means the National Commission on State Workmen's Compensation Laws established under this Act.
SEC.
4.
Applicability of This Act
(a)

This Act shall apply with respect to employment performed in a workplace in a State, the District of Columbia, the Commonwealth of Puerto Rico, the Virgin Islands, American Samoa, Guam, the Trust Territory of the Pacific Islands, Wake Island, Outer Continental Shelf Lands defined in the Outer Continental Shelf Lands Act, Johnston Island, and the Canal Zone. The Secretary of the Interior shall, by regulation, provide for judicial enforcement of this Act by the courts established for areas in which there are no United States district courts having jurisdiction.
(b)
 
(1)

Nothing in this Act shall apply to working conditions of employees with respect to which other Federal agencies, and State agencies acting under section 274 of the Atomic Energy Act of 1954, as amended (42 U.S.C. 2021), exercise statutory authority to prescribe or enforce standards or regulations affecting occupational safety or health.
(2)

The safety and health standards promulgated under the Act of June 30, 1936, commonly known as the Walsh-Healey Act (41 U.S.C. 35 et seq.), the Service Contract Act of 1965 (41 U.S.C. 351 et seq.), Public Law 91-54, Act of August 9, 1969 (40 U.S.C. 333), Public Law 85-742, Act of August 23, 1958 (33 U.S.C. 941), and the National Foundation on Arts and Humanities Act (20 U.S.C. 951 et seq.) are superseded on the effective date of corresponding standards, promulgated under this Act, which are determined by the Secretary to be more effective. Standards issued under the laws listed in this paragraph and in effect on or after the effective date of this Act shall be deemed to be occupational safety and health standards issued under this Act, as well as under such other Acts.
(3)

The Secretary shall, within three years after the effective date of this Act, report to the Congress his recommendations for legislation to avoid unnecessary duplication and to achieve coordination between this Act and other Federal laws.
(4)

Nothing in this Act shall be construed to supersede or in any manner affect any workmen's compensation law or to enlarge or diminish or affect in any other manner the common law or statutory rights, duties, or liabilities of employers and employees under any law with respect to injuries, diseases, or death of employees arising out of, or in the course of, employment.

"""
prompt = (
    "Below is an instruction that describes a task, paired with an input that "
    "provides further context. Write a response that appropriately completes the request.\n\n"
    "### Instruction:\n"
    f"{instruction}\n\n"
    "### Input:\n"
    f"{input_txt}\n\n"
    "### Response:\n"
)
tokens = tok(prompt, return_tensors="pt").to(model.device)

out = model.generate(
    **tokens,
    max_new_tokens=256,
    temperature=0.7,
    top_p=0.9,
)
print(tok.decode(out[0], skip_special_tokens=True).split("### Response:")[-1].strip())

Output

9. Conclusion and Recap

This blog has detailed a methodology for fine-tuning the Qwen-2.5–14B-Instruct model using Low-Rank Adaptation (LoRA) on an 8x NVIDIA A100–40GB GPU system, specifically focusing on handling long input sequences (4096 tokens). The approach leverages the Axolotl framework to orchestrate a distributed training strategy combining Distributed Data Parallelism (DDP) and Sequence Parallelism (SP).

The core challenge addressed is the significant memory requirement imposed by long sequences within the Transformer attention mechanism. This was tackled through a multi-pronged strategy configured via Axolotl:

1. Parameter Efficiency: Using LoRA (adapter: lora, lora_r: 16) minimized the number of trainable parameters and associated optimizer state memory.
2. Memory-Saving Techniques: Employing 8-bit base model loading (load_in_8bit: True), bfloat16 mixed precision (bf16: True), gradient checkpointing (gradient_checkpointing: True), and an 8-bit optimizer (optimizer: adamw_bnb_8bit) substantially reduced the memory footprint.
3. Optimized Attention: Utilizing FlashAttention (flash_attention: True) provided memory and speed benefits crucial for the 4096 sequence length.
4. Sequence Parallelism: Distributing the sequence computation across GPUs (sequence_parallel_degree: 4) using ring-flash-attn kernels (ring_attn_func: "batch_ring") directly reduced the per-GPU activation memory related to sequence length.
5. Data Parallelism: Using DDP (managed by Accelerate) parallelized the processing of data batches across all 8 GPUs, increasing overall throughput.

The combination of DDP and SP proved effective for this LoRA fine-tuning scenario. DDP ensures data processing scales across GPUs, while SP specifically targets the sequence length memory bottleneck, which LoRA itself does not address. This contrasts with strategies like FSDP or TP, which might be less optimal here due to LoRA’s inherent parameter efficiency (reducing the need for parameter/optimizer sharding) and the primary challenge being sequence length, not individual layer size.

Axolotl played a key role in simplifying the implementation of this complex workflow. Its YAML-based configuration allowed for straightforward specification of the base model, LoRA parameters, dataset, optimization settings, precision, and the combined DDP+SP distributed strategy, including the necessary ring-flash-attn kernels available in version 0.8.1.

Links:

1. Github Repo: https://github.com/dhnanjay/ft_using_Axolotl_1