Llama-3.1-8B-Instruct-w16a8-1node-bs8

Abstract

The Llama-3.1-8B-Instruct-w16a8-1node-bs8 model is a Turkish legal instruction-tuned variant of Llama-3.1-8B-Instruct, trained using Float8 mixed precision with the default Tensorwise quantization scaling recipe under the FSDP2 distributed training framework. This model represents the single-node FP8 configuration in the “FSDP2 with Float8 Precision for Faster Training” experiment, designed to evaluate the impact of FP8 activations and BF16 weights on training speed, memory efficiency, and legal-domain performance.

Experiment Context

This model was trained as part of our study for comparing FSDP2 with bfloat16 precision against FSDP2 with FP8 mixed precision bfp16-fp8. We used meta-llama/Llama-3.1-8B-Instruct. The model has been loaded using torch_dtype = bfloat16 and for FP8 + FSDP2 compatibility the model has been wrap per-layer instead of whole model This helped to avoid dimension misalignment issues and during forward and backward passes float8 variats been used for computations where FP8E4M3 for activations (forward pass) and FP8E5M2 for gradients (backward pass, wider range) also we setted the pad_inner_dim for automatically pad dimensions to be divisible by 16 which is required for FP8.

from torchao.float8 import (
    convert_to_float8_training,
    Float8LinearConfig,
    precompute_float8_dynamic_scale_for_fsdp)
config = Float8LinearConfig(
  pad_inner_dim=True,
  enable_fsdp_float8_all_gather=True)
model = convert_to_float8_training(model, config=config)
if use_fp8:
    for i, layer in enumerate(model.model.layers):
        fully_shard(layer, **fsdp_kwargs)
    fully_shard(model.model.embed_tokens, **fsdp_kwargs)
    fully_shard(model.lm_head, **fsdp_kwargs)

Base Model Technical Specifications

• Parameters: 8 Billion
• Architecture Family: Llama 3.1
• Maximum Position Embeddings: 131,072
• Attention Heads: 32 (num_attention_heads)
• Key-Value Heads: 8 (num_key_value_heads)
• Hidden Layers: 32 (num_hidden_layers)
• Hidden Size: 4,096 (hidden_size)
• Intermediate Size: 14,336
• Vocabulary Size: 128,256
• Precision: bfloat16
• RoPE Scaling: type llama3, factor = 8.0
• RMS Norm Epsilon: 1e-05
• Activation: SiLU

Training Methodology

Training Configuration

• Model: meta-llama/Llama-3.1-8B-Instruct
• Sequence Length: 4,096 (seq_len)
• Epochs: 2
• Per-Device Micro Batch Size: 2
• Gradient Accumulation: 4
• GPUs: 4 (via CUDA_VISIBLE_DEVICES=0,1,2,3)
• dtype: bf16 && fp8=true
  • Weights: bfloat16
  • Activations: float8
• Optimizer: AdamW
  • Learning Rate: 2e-5
  • Weight Decay: 0.01
  • Betas: (0.9, 0.95)
  • Epsilon: 1e-8
• LR Scheduler: Cosine; warmup = 10% (warmup_ratio=0.1) | also warmup_steps=100
• Max Grad Norm: 1.0
• Gradient Checkpointing: Enabled
• Checkpointing: every 10 steps; keep last 5; select best by eval_loss
• Logging: every step to file; Weights & Biases in offline mode
• Seed: 100
• Distributed Training: torch.distributed.run (single node, multi-GPU)
  • FSDP2 (Optimized Fully Sharded Data Parallel)

Setups

• Precision: Used Half-precision bfloat16 as data type and for computation.
• Hardware: HPC (EuroHPC/BSC-class) node with 4 × NVIDIA H100 GPUs.
• Framework: PyTorch with torchrun for distributed training.

Dependencies

Job Details

All 15-models trained on(1Node,4Noes,8Nodes with both bfp16-fp8 && bfp16 configurations and fp8 recipes)

Implementation

Gpu && Memory usage Profiling

To visualize the usage of the memory and gpu's vram/utilis and more you can profile them with pytorch-profiler tool or nsight systems profiler tool .

• follow the steps to visualize with Pytorch-profiler :

  1. pip install the versions that mentioned in the dependencies section of these libs tensorboard and tensorboard-data-server.
  2. Visualize pytorch profiles by runing the command provided below.

  tensorboard --logdir="/gpfs/scratch/ehpc317/esavirdi/dual_precision_training/profiles/Llama-3.1-8B-Instruct_tensorwise" --port="6006"
  

• follow the steps to visualize with nsight systems:

  1. download the files on to you local machine.
  2. download the full_version of the nsys profiler tool from here : https://developer.nvidia.com/nsight-systems/get-started.
  3. open the nsys-rep profiles using the nsight systems gui .

  also you can get a summary of the report using these command lines up to your porpuse :

  nsys stats --report cuda_gpu_kern_sum /path/to/lama3.1_fp16_fp8.nsys-rep
  

  nsys stats --report cuda_gpu_trace /path/to/lama3.1_fp16_fp8.nsys-rep
  

    nsys stats --report cuda_gpu_mem_size_sum /path/to/lama3.1_fp16_fp8.nsys-rep
  

Usage

Note: the final model has saved in bfloat16 format. For inference, load the model in bfloat16 or float16 as shown below:

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
model_name = "newmindai/Llama-3.1-8B-Instruct-w16a8-1node-bs8"
dtype = torch.bfloat16
tok = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=dtype,
    device_map="auto"
)

prompt = "Soru: Kişisel Verilerin Korunması Kanunu uyarınca hangi durumlarda açık rıza aranmaz? Cevap:"
inputs = tok(prompt, return_tensors="pt").to(model.device)
with torch.no_grad():
    out = model.generate(
        **inputs,
        max_new_tokens=256,
        do_sample=False
    )

print(tok.decode(out[0], skip_special_tokens=True))

Ethical Considerations and Disclaimers

• Research & development purposes only; not a substitute for professional legal counsel.
• Users must ensure compliance with data protection and sector regulations.
• Potential biases may exist in domain data and model outputs.

Model & Data Card Metadata

• Total Parameters: 8,030,261,248
• Serialized Size (approx.): 16,060,522,496 bytes
• Config precision: bfloat16
• RoPE: llama3 scaling, factor 8.0

References and Citations

Base Model

@misc{meta_llama31_8b_instruct,
  title={Llama 3.1 8B Instruct},
  author={Meta AI},
  year={2024},
  howpublished={\url{https://huggingface.co/meta-llama/Llama-3.1-8B-Instruct}}
}

Training Dataset

@misc{euro_hpc_legal,
  title={EuroHPC-Legal},
  author={newmindai},
  year={2025},
  howpublished={\url{https://huggingface.co/datasets/newmindai/EuroHPC-Legal}}
}