Fine-Tuning Llama 3.1 405B on NVIDIA B200 GPUs

Llama 3.1 405-billion (or 405B) from Meta is one of the first frontier-level open-source AI models that rivals the performance of other closed-source models such as Claude 3 and GPT-4. Llama 3.1 405B is a 405B-parameter dense autoregressive transformer model with a supported 128K-token context window.

Fine tuning a 405B parameter model requires planning related to GPU memory allocation, infrastructure costs, and distributed training frameworks. In this blog, we deep dive into the mechanics of fine tuning Llama 3.1 405B on NVIDIA B200 GPUs that feature 180GB of HBM3e memory per GPU.

B200-class systems provide strong capacity for inference and serving with TensorRT-LLM, including FP8/FP4 paths where supported. For fine-tuning, the core stack is usually PyTorch with distributed training frameworks such as NeMo/Megatron-LM, DeepSpeed ZeRO-3, FSDP or carefully configured PEFT/QLoRA; TensorRT-LLM should not be positioned as the fine-tuning framework.

Overview Of Llama 3.1 405B

Llama 3.1 405B, released in July 2024, is considered as one of the largest models of the Llama 3.1 family. It is a dense, auto-regressive transformer model and one of the first open-source models to reach frontier-level performance. Hence, all the 405 billion parameters participate during the process of inference and training.

Here are some of the salient features of Llama 3.1 405B model:

Apart from 405B (frontier-level), the Llama 3.1 model was also released in two more sizes – 8B (lightweight) and 70B (balance of efficiency and performance). Here are some of major capabilities of the Llama 3.1 405B model:

As per the competitive benchmarking results, Llama 3.1 405B performed very closely to GPT-4o and Claude 3.5 Sonnet on reasoning and knowledge benchmarks:

Llama 3.1 405B is built for usage in modern agentic enterprise AI workflows, as it is designed to interact with external tools, APIs, databases, etc. for generating structured outputs.

NVIDIA B200 Architecture For AI Workloads

Based on the Blackwell architecture, NVIDIA B200 GPU is primarily built for efficient handling of next-generation AI and high-performance computing workloads. A single B200 GPU includes a dual-die configuration where two GPU dies comprise 208 billion transistors.

B200 features the following enhancements:

• 4 L2 cache partitions (double to those in Hopper)
• 8 HBM3e memory stack
• 192 GB of HBM3e memory
• Support for advanced low-precision computation formats – FP4, FP8, BF16
• Improved Tensor cores help accelerate matrix multiplications, mixed-precision computations, and deep-learning tensor operations
• High-speed NVLink/NVSwitch interconnects help multi-GPU communication within systems such as DGX B200, but multi-node training still depends on the external network fabric, storage path and distributed training configuration

For Llama 3.1 405B workloads, B200 GPUs provide the memory bandwidth, compute throughput, and scalability required for large-context inference, distributed model serving, parameter-efficient fine-tuning, and multi-node model sharding. Let’s deep dive into it.

In comparison to 80 GB HBM2e memory on A100 and up to 94 GB HBM3 on H100 GPUs, the B200 can deliver up to 192 GB HBM3e/GPU. What this essentially means is that the B200 GPU will need fewer GPUs to fit the Llama 3.1 405B model, which directly cuts costs per run!

Fine-tuning Llama 3.1 405B can be constrained by memory capacity, activation memory, memory bandwidth, data loading, checkpointing, interconnect, and optimizer state. The training performance depends on the efficiency of data movement between high-bandwidth GPU memory and compute cores.

The 8 TB/s transfer speed offered by B200 GPU results in faster model weight reads per forward/backward pass. This in turn reduces the time taken per training step.

All these factors associated with the B200 architecture help improve the overall training throughput and scalability for frontier-scale models such as Llama 3.1 405B.

Memory Anatomy Of A 405B Fine-Tuning Run

BF16 [bfloat16] is a preferred precision format for large-scale LLM training and fine-tuning. As stated earlier, fine-tuning a large Llama 3.1 405B model is a memory-intensive workload. Like any other model, the model weights require a good amount of GPU memory capacity.

At BF16 precision, each parameter occupies 2 bytes. Hence, a 405B parameter model requires close to 405B * 2 bytes ≈ 810 GB of memory. Model weights is just one aspect of fine-tuning, factors like gradients, optimizer states, activations, and distributed communication buffers also contribute to the consumption of GPU memory.

As mentioned earlier, a model stored in BF16 precision may require 2 bytes per parameter for the weights. However, copies of data are often needed for gradients and optimizer calculations during training. This increases the overall memory requirement. Owing to this, a full-parameter fine-tuning of Llama 3.1 405B can require multiple terabytes of aggregate GPU memory.

Multiple B200-class GPUs can be used for 405B fine-tuning only with proper distributed training techniques such as FSDP, DeepSpeed ZeRO-3, tensor/pipeline parallelism, or NeMo/Megatron-LM.

Approaches For Fine-Tuning Llama 3.1 405B

Like any other model, there are a number of different approaches for fine-tuning the Llama 3.1 405B. This includes full-fine tuning (or full-parameter fine tuning) to more memory-efficient techniques such as Low-Rank Adaptation (LoRA), Quantized Low-Rank Adaptation (QLoRA), and Parameter-Efficient Fine-Tuning (PEFT).

Full-Parameter Fine Tuning

As the name suggests, full-fine tuning involves updating all the parameters of the Llama 3.1 405B model during the training process. It is the least memory- and parameter-efficient approach compared with LoRA/QLoRA/PEFT, and it is usually not the first practical choice unless the adaptation requirement justifies the cost.

The overall memory footprint for a 405B model is large considering that GPU memory is consumed by model weights, gradients, activations, optimizer states, and more.

Here is the approximate amount of GPU memory required for just storing the model weights alone in the BF16 precision:

• 405 × 10⁹ → 405 billion parameters in Llama 3.1 405B
• 2 bytes → 2 bytes of memory is occupied by every BF16 parameter

This is ≈ 810 GB (or 810 billion bytes) for simply storing the model weights. Hence, a single NVIDIA B200 GPU with 192 GB of HBM3e memory (per GPU) might still not be sufficient for full-parameter fine-tuning of a 405B model. Multiple NVIDIA B200 GPUs working together in a distributed setup can be instrumental in making large-scale fine-tuning feasible.

Low-Rank Adaptation (LoRA)

Low-Rank Adaptation, commonly known as LoRA, is a popular PEFT technique where the base model weights are frozen and small trainable adapter matrices are added to selected layers.

While the model weight is frozen, only a fraction of the total weights are used for training. As fewer parameters are being trained, LoRA requires significantly less GPU memory than full fine-tuning. This directly impacts the GPU and model training costs. The lighter compute footprint also makes it practical to run multiple experiments quickly.

On completion of the training, the saved checkpoints contain only the adapter weights and not the full copy of the model. Full fine-tuning of Llama 3.1 405B in BF16 demands upward of 810 GB just for weights alone, LoRA’s reduced parameter footprint can cut the required GPU count by a huge margin.

Quantized Low-Rank Adaptation (QLoRA)

Quantized Low-Rank Adaptation, commonly known as QLoRA, takes the benefits of LoRA way further by combining LoRA with quantization. The foundational model is initialized using 4-bit precision, a technique that significantly reduces the overall memory footprint of the weights.

The LoRA adapters are still trained in higher precision on top of it. By utilizing QLoRA with 4-bit NormalFloat (NF4) quantization, the foundational model weight footprint is substantially minimized to approximately 202 GB for the 405B-parameter model.

• 405 × 10⁹ x 0.5 bytes ≈ 202.5 GB

This reduction in memory footprint can be instrumental in large-scale fine-tuning of 405B-parameter models with smaller GPU setup.

The choice of fine-tuning the Llama 3.1 405B depends on the GPU infrastructure, training budget, and adaptation required for the target use case.

Impact Of Batch Size On Training Performance

There is a significant impact of batch size when fine-tuning frontier-scale models such as Llama 3.1 405B. Though larger batch size results in an increased throughput, it does increase the activation memory which in turn can become one of the biggest VRAM consumers.

The aforementioned dynamics are especially critical when considering a frontier-level model of the 405B parameter scale. With up to 192 GB of HBM3e memory and ~8 TB/s memory bandwidth, NVIDIA B200 GPU helps in running larger batch sizes without hitting memory ceilings.

This can lower the total GPU requirement and improve overall cost efficiency per training run.

GPU Memory Utilization During Training

Memory usage fluctuates across forward pass, backward pass, optimizer step, activations, gradients, and optimizer states.

High GPU utilization is desirable for attaining maximum efficiency but it is not ideal to run too close to full VRAM. This could lead to Out-Of-Memory (OOM) errors due to the dynamic memory spikes during the training process.

Target high but safe memory utilization only after profiling. For 405B fine-tuning, keep headroom for activation spikes, communication buffers, checkpointing, dataloader variance and framework fragmentation. GPU memory usage does not remain static during the training cycles.

Activations are stored during the forward pass and gradients are computed during the backward pass. For full fine-tuning, Adam/AdamW optimizer states can dominate memory. For LoRA/QLoRA, optimizer state applies mainly to adapter parameters, but activations and frozen-base sharding still remain major memory consumers.

Figure 1: GPU memory usage during training cycles [Image Source]

Here are the two techniques that improve the utilization of 405B workloads:

• For models that as large as 405B with support for context lengths up to 128K tokens, Sequence packing helps eliminate the padding waste by combining multiple shorter sequences into a single training example. This increases token throughput for variable-length instruction-tuning datasets.
• Gradient checkpointing is a technique used to significantly reduce VRAM usage during training. For frontier-scale models like Llama 3.1 405B, it reduces activation memory by recomputing intermediate activations during the backward pass instead of storing them.

Practical Llama 3.1 405B Fine-Tuning Implementation with PyTorch, LoRA, and NVIDIA B200

As discussed earlier, full fine-tuning of Llama 3.1 405B can require multi-terabyte aggregate memory, so parameter-efficient approaches such as LoRA or QLoRA are usually the best choices for practical implementation.

Since LoRA also reduces memory consumption and model training costs, it is extensively used in production environments. The example shown below demonstrates a LoRA-based fine-tuning pipeline using PyTorch and Hugging Face PEFT (Parameter-Efficient Fine-Tuning).

The local Apple Silicon example can validate dataset formatting and adapter-training flow on a small model, but it cannot prove that the same script will scale to Llama 3.1 405B on distributed B200 GPUs. Scaling requires a separate distributed training configuration, memory plan and framework validation.

Here is a gist of what all is covered in the example implementation:

• Loading a pretrained Llama-compatible model
• Applying LoRA adapters to transformer attention layers
• Preparing an instruction-style JSONL dataset
• Tokenizing and formatting prompts
• Running supervised fine-tuning using Hugging Face Trainer
• Saving LoRA adapter checkpoints
• Running inference using the fine-tuned adapters

In order to run the code, you need to have an account on Hugging Face, post which you have to generate a user access token by navigating to the Hugging Face Tokens section. Once the token is generated, make sure to copy the token since it will be used in the latter section of the implementation.

Figure 2: Hugging Face Tokens [Image Source]

The implementation also showcases why LoRA is extensively used for large models such as Llama 3.1 405B. Let’s look at the implementation:

Step 1 – Install Dependencies

Since the model validation is done on Apple Silicon, we create two separate requirements files:

• requirements-mac.txt – Used for validation of smaller model (i.e. Llama-3.2-1B on Apple Silicon)
• requirements-b200.txt – Used for validation of Llama 3.1 405B on NVIDIA B200 GPUs on AceCloud

Shown below are the dependencies mentioned in requirements-mac.txt:

torch>=2.4.0
transformers>=4.44.0
peft>=0.12.0
accelerate>=0.33.0
datasets>=2.21.0
trl>=0.10.1
tensorboard>=2.17.0
sentencepiece>=0.2.0
scipy>=1.14.0
protobuf>=5.27.0

To get started, create a venv by running the commands python -m venv venv and source venv/bin/activate on the terminal. Once the venv is created, install the dependencies by triggering the command pip install -r requirements-mac.txt on the terminal.

PyTorch, Transformers, PEFT, Datasets, and Accelerate are some of the notable dependencies installed in the fine-tuning environment for training.

Here the requirements file (i.e., requirements-b200.txt) needs to be installed on the GPU environment (or VM) where the actual fine-tuning and inference workloads will run.

torch>=2.4.0
torchvision>=0.19.0
torchaudio>=2.4.0
transformers>=4.44.0
peft>=0.12.0
bitsandbytes>=0.43.3
accelerate>=0.33.0
datasets>=2.21.0
trl>=0.10.1
flash-attn>=2.6.3
tensorboard>=2.17.0
scipy>=1.14.0
sentencepiece>=0.2.0
protobuf>=5.27.0

Step 2 – Configure the Base Model

For local validation, smaller models such as TinyLlama or Llama 3.2 1B can be used before scaling the same workflow to Llama 3.1 405B on NVIDIA B200 available in the GPU Cloud.

In this example, we have used Llama 3.2 1B. For using the model, you first need to raise a request to access model meta-llama/Llama-3.2-1. Post acceptance (which normally takes ~5 minutes), you would receive an email from Hugging Face confirming the access to the gated Llama 3.1 model repository.

Figure 3: Email from Hugging Face

Trigger the command hf auth login to login to Hugging Face using the generated token.

Step 3 – Prepare the Dataset

A JSONL instruction-style dataset is used, where each training sample contains an instruction and its corresponding output response. A sample .jsonl file is shown below:

{"instruction":"Explain LoRA fine-tuning","output":"LoRA fine-tuning trains only a small subset of parameters instead of the full model."}
{"instruction":"What is parameter-efficient fine-tuning?","output":"Parameter-efficient fine-tuning updates a small number of trainable weights while keeping the base model frozen."}
{"instruction":"Why is LoRA useful?","output":"LoRA significantly reduces GPU memory consumption during large language model fine-tuning."}
{"instruction":"What are LoRA adapters?","output":"LoRA adapters are lightweight trainable matrices injected into transformer layers."}

FileName – dataset/train.jsonl

In preprocessing, instruction-response pairs should be converted using the tokenizer’s official chat template where possible, with labels masked so loss is applied mainly to the assistant response rather than forcing the model to learn the user prompt.

Step 4 – Apply LoRA Adapters

Rather than updating the entire 405-billion parameter set during the adaptation process, LoRA freezes the foundational model weights and integrates lightweight, trainable adapter layers within key transformer components, including:

• q_proj – Query Projection
• k_proj – Key Projection
• v_proj – Value Projection
• o_proj – Output Projection

Here is a code snippet of the LoRA configuration:

lora_config = LoraConfig(
    r=LORA_RANK,
    lora_alpha=LORA_ALPHA,
    target_modules=[
        "q_proj",
        "k_proj",
        "v_proj",
        "o_proj",
    ],
    lora_dropout=LORA_DROPOUT,
    bias="none",
    task_type=TaskType.CAUSAL_LM,
)

FileName – fine-tuning/qlora_finetune_llama405b_mac.py

Step 5 – Run Fine-Tuning

For fine-tuning the Llama-3.2-1B on Apple Silicon, we use the Hugging Face Trainer API along with small batch sizes, gradient checkpointing, and other memory-efficient training optimizations.

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset,
    processing_class=tokenizer,
    data_collator=DataCollatorForSeq2Seq(
        tokenizer,
        pad_to_multiple_of=8,
        return_tensors="pt",
        padding=True,
    ),
)

FileName – fine-tuning/qlora_finetune_llama405b_mac.py

The same example can also be extended to large-scale GPU infrastructure such as NVIDIA B200 on AceCloud for fine-tuning models like Llama 3.1 405B. However, this would require changes related to distributed sharding, multi-node orchestration, checkpoint strategy and validated memory planning.

For simplification, we have changed the model definition as shown below:

MODEL_NAME = "meta-llama/Meta-Llama-3.1-405B"
DATASET_PATH = "dataset/train.jsonl"
OUTPUT_DIR = "./llama31-405b-qlora"

FileName – fine-tuning/qlora_finetune_llama405b.py

Here are some of the additional changes for fine-tuning Llama 3.1 405B on NVIDIA B200 GPUs:

• Enabling multi-GPU distributed training across NVIDIA B200 GPUs
• Using 4-bit NF4 quantization for QLoRA fine-tuning
• Increasing gradient accumulation steps to support larger effective batch sizes
• Enabling FlashAttention for faster and more memory-efficient attention computation
• Applying gradient checkpointing to reduce activation memory usage
• Using sequence packing to minimize padding overhead
• Configuring BF16 precision for optimized training performance
• Adjusting LoRA adapter configurations for large-scale transformer layers
• Tensor parallelism / FSDP / DeepSpeed

Here is the complete implementation of fine-tuning Llama 3.1 405B model on NVIDIA B200 GPUs:

import os
import torch

from datasets import load_dataset
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    TrainingArguments,
    Trainer,
    DataCollatorForSeq2Seq,
)

from peft import (
    LoraConfig,
    get_peft_model,
    prepare_model_for_kbit_training,
    TaskType,
)

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
MODEL_NAME = "meta-llama/Meta-Llama-3.1-405B"
DATASET_PATH = "dataset/train.jsonl"
OUTPUT_DIR = "./llama31-405b-qlora"

MAX_SEQ_LENGTH = 4096
LORA_RANK = 16
LORA_ALPHA = 32
LORA_DROPOUT = 0.05

# ---------------------------------------------------------------------------
# 4-bit NF4 quantization config
# ---------------------------------------------------------------------------
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,
)

# ---------------------------------------------------------------------------
# Tokenizer
# ---------------------------------------------------------------------------
tokenizer = AutoTokenizer.from_pretrained(
    MODEL_NAME,
    use_fast=True,
)

tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"

# ---------------------------------------------------------------------------
# Model
# ---------------------------------------------------------------------------
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    quantization_config=bnb_config,
    device_map="auto",
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2",
)

# Prepare quantized model for LoRA fine-tuning
model = prepare_model_for_kbit_training(
    model,
    use_gradient_checkpointing=True,
)

# ---------------------------------------------------------------------------
# LoRA configuration
# ---------------------------------------------------------------------------
lora_config = LoraConfig(
    r=LORA_RANK,
    lora_alpha=LORA_ALPHA,
    target_modules=[
        "q_proj",
        "k_proj",
        "v_proj",
        "o_proj",
        "gate_proj",
        "up_proj",
        "down_proj",
    ],
    lora_dropout=LORA_DROPOUT,
    bias="none",
    task_type=TaskType.CAUSAL_LM,
)

model = get_peft_model(model, lora_config)
model.print_trainable_parameters()

# ---------------------------------------------------------------------------
# Dataset
# ---------------------------------------------------------------------------
# Expected JSONL format:
# {"instruction": "...", "output": "..."}
dataset = load_dataset(
    "json",
    data_files={"train": DATASET_PATH},
    split="train",
)

def format_prompt(sample):
    return (
        f"<|begin_of_text|>"
        f"<|start_header_id|>user<|end_header_id|>\n"
        f"{sample['instruction']}"
        f"<|eot_id|>"
        f"<|start_header_id|>assistant<|end_header_id|>\n"
        f"{sample['output']}"
        f"<|eot_id|>"
    )

def tokenize(sample):
    prompt = format_prompt(sample)
    result = tokenizer(
        prompt,
        truncation=True,
        max_length=MAX_SEQ_LENGTH,
        padding=False,
    )
    result["labels"] = result["input_ids"].copy()
    return result

tokenized_dataset = dataset.map(
    tokenize,
    remove_columns=dataset.column_names,
    desc="Tokenizing dataset",
)

# ---------------------------------------------------------------------------
# Training arguments
# ---------------------------------------------------------------------------
training_args = TrainingArguments(
    output_dir=OUTPUT_DIR,
    per_device_train_batch_size=2,
    gradient_accumulation_steps=16,
    bf16=True,
    fp16=False,
    optim="paged_adamw_32bit",
    learning_rate=2e-4,
    lr_scheduler_type="cosine",
    warmup_ratio=0.03,
    weight_decay=0.01,
    max_grad_norm=1.0,
    num_train_epochs=3,
    save_strategy="steps",
    save_steps=500,
    save_total_limit=3,
    logging_steps=10,
    report_to="tensorboard",
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"use_reentrant": False},
    group_by_length=True,
    dataloader_num_workers=4,
)

# ---------------------------------------------------------------------------
# Trainer
# ---------------------------------------------------------------------------
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset,
    tokenizer=tokenizer,
    data_collator=DataCollatorForSeq2Seq(
        tokenizer,
        pad_to_multiple_of=8,
        return_tensors="pt",
        padding=True,
    ),
)

# ---------------------------------------------------------------------------
# Train
# ---------------------------------------------------------------------------
trainer.train()

# ---------------------------------------------------------------------------
# Save adapter weights and tokenizer
# ---------------------------------------------------------------------------
model.save_pretrained(OUTPUT_DIR)
tokenizer.save_pretrained(OUTPUT_DIR)

print(f"Adapter weights saved to {OUTPUT_DIR}")

FileName – fine-tuning/qlora_finetune_llama405b.py

Here is the complete implementation of the LoRA fine-tuning pipeline that will be done on Apple Silicon (m4 pro):

import os
import torch

from datasets import load_dataset
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    TrainingArguments,
    Trainer,
    DataCollatorForSeq2Seq,
)

from peft import (
    LoraConfig,
    get_peft_model,
    prepare_model_for_kbit_training,
    TaskType,
)

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
MODEL_NAME = "meta-llama/Meta-Llama-3.1-405B"
DATASET_PATH = "dataset/train.jsonl"
OUTPUT_DIR = "./llama31-405b-qlora"

MAX_SEQ_LENGTH = 4096
LORA_RANK = 16
LORA_ALPHA = 32
LORA_DROPOUT = 0.05

# ---------------------------------------------------------------------------
# 4-bit NF4 quantization config
# ---------------------------------------------------------------------------
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,
)

# ---------------------------------------------------------------------------
# Tokenizer
# ---------------------------------------------------------------------------
tokenizer = AutoTokenizer.from_pretrained(
    MODEL_NAME,
    use_fast=True,
)

tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"

# ---------------------------------------------------------------------------
# Model
# ---------------------------------------------------------------------------
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    quantization_config=bnb_config,
    device_map="auto",
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2",
)

# Prepare quantized model for LoRA fine-tuning
model = prepare_model_for_kbit_training(
    model,
    use_gradient_checkpointing=True,
)

# ---------------------------------------------------------------------------
# LoRA configuration
# ---------------------------------------------------------------------------
lora_config = LoraConfig(
    r=LORA_RANK,
    lora_alpha=LORA_ALPHA,
    target_modules=[
        "q_proj",
        "k_proj",
        "v_proj",
        "o_proj",
        "gate_proj",
        "up_proj",
        "down_proj",
    ],
    lora_dropout=LORA_DROPOUT,
    bias="none",
    task_type=TaskType.CAUSAL_LM,
)

model = get_peft_model(model, lora_config)
model.print_trainable_parameters()

# ---------------------------------------------------------------------------
# Dataset
# ---------------------------------------------------------------------------
# Expected JSONL format:
# {"instruction": "...", "output": "..."}
dataset = load_dataset(
    "json",
    data_files={"train": DATASET_PATH},
    split="train",
)

def format_prompt(sample):
    return (
        f"<|begin_of_text|>"
        f"<|start_header_id|>user<|end_header_id|>\n"
        f"{sample['instruction']}"
        f"<|eot_id|>"
        f"<|start_header_id|>assistant<|end_header_id|>\n"
        f"{sample['output']}"
        f"<|eot_id|>"
    )

def tokenize(sample):
    prompt = format_prompt(sample)
    result = tokenizer(
        prompt,
        truncation=True,
        max_length=MAX_SEQ_LENGTH,
        padding=False,
    )
    result["labels"] = result["input_ids"].copy()
    return result

tokenized_dataset = dataset.map(
    tokenize,
    remove_columns=dataset.column_names,
    desc="Tokenizing dataset",
)

# ---------------------------------------------------------------------------
# Training arguments
# ---------------------------------------------------------------------------
training_args = TrainingArguments(
    output_dir=OUTPUT_DIR,

    # Effective batch size:
    # per_device_train_batch_size
    # x gradient_accumulation_steps
    # x num_gpus
    per_device_train_batch_size=2,
    gradient_accumulation_steps=16,

    bf16=True,
    fp16=False,
    optim="paged_adamw_32bit",
    learning_rate=2e-4,
    lr_scheduler_type="cosine",
    warmup_ratio=0.03,
    weight_decay=0.01,
    max_grad_norm=1.0,
    num_train_epochs=3,
    save_strategy="steps",
    save_steps=500,
    save_total_limit=3,
    logging_steps=10,
    report_to="tensorboard",
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"use_reentrant": False},
    group_by_length=True,
    dataloader_num_workers=4,
)

# ---------------------------------------------------------------------------
# Trainer
# ---------------------------------------------------------------------------
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset,
    tokenizer=tokenizer,
    data_collator=DataCollatorForSeq2Seq(
        tokenizer,
        pad_to_multiple_of=8,
        return_tensors="pt",
        padding=True,
    ),
)

# ---------------------------------------------------------------------------
# Train
# ---------------------------------------------------------------------------
trainer.train()

# ---------------------------------------------------------------------------
# Save adapter weights and tokenizer
# ---------------------------------------------------------------------------
model.save_pretrained(OUTPUT_DIR)
tokenizer.save_pretrained(OUTPUT_DIR)

print(f"Adapter weights saved to {OUTPUT_DIR}")

FileName – fine-tuning/qlora_finetune_llama405b_mac.py

Trigger the command python fine-tuning/qlora_finetune_llama405b_mac.py to run the fine-tuning on the Apple Silicon:

As seen from the execution snapshot, the fine-tuned LoRA adapter weights are successfully saved in the ./llama32-1b-lora directory.

Along with the command torchrun --nproc_per_node=8 fine-tuning/qlora_finetune_llama405b.py, you must also provide an Accelerate, DeepSpeed, FSDP, or NeMo launch configuration.

The configuration should define tensor, sequence, and data parallelism, along with node count, sharding policy, mixed precision settings, checkpoint paths, and NCCL/network configurations required for B200-scale distributed environments.

Step 6 – Run Inference (Optional)

After training, the LoRA adapter weights can be reloaded alongside the base model to generate responses using the fine-tuned parameters. This is used for validating the successful execution of the fine-tuning pipeline before it is scaled to larger B200 GPU clusters.

The inference output only confirms that the small-model/local adapter path works. It still needs to be checked thoroughly if Llama 3.1 405B fine-tuning or inference runs successfully on B200 infrastructure.

Best Practices For Large-Scale LLaMA 3.1 Fine-Tuning

As seen throughout the course of the blog, large-scale fine-tuning of a large model like Llama3.1 405B requires optimizations across areas like memory utilization, batch sizing, distributed training, and GPU communication.

Fine-tuning techniques like LoRA, in conjunction with gradient checkpointing and quantization can help reduce GPU infrastructure costs without compromising on the scalability.

NVIDIA B200 GPUs on AceCloud can be leveraged for Llama 3.1 405B experimentation and parameter-efficient fine-tuning. This should be done only after thorough validation of GPU availability, exact GPU memory exposed, distributed framework support, storage throughput, interconnect, cost per run, and support scope.

Author Popular Post

Jason Karlin

author

Industry veteran with over 10 years of experience architecting and managing GPU-powered cloud solutions. Specializes in enabling scalable AI/ML and HPC workloads for enterprise and research applications. Former lead solutions architect for top-tier cloud providers and startups in the AI infrastructure space.

• CUDA cores vs Tensor cores: Choosing the Right GPU for Machine Learning  
• GPUs in Healthcare: Transforming Diagnoses and Treatment with Cutting-Edge Technology
• The New Wave of Cloud GPUs: Revolutionizing the Business Landscape
• Harnessing GPU-as-a-Service: Key Benefits and Future Trends 
• Best GPUs For Deep Learning [2025 Updated List]