Llama-3.1-8B-Instruct-w16a8-rw-4nodes

Abstract

The Llama-3.1-8B-Instruct-w16a8-rw model is a domain-adapted Turkish legal instruction-tuned variant of Meta’s Llama-3.1-8B-Instruct, trained using the Float8 Rowwise quantization recipe. This model was developed within the “FSDP2 with Float8 Precision for Faster Training” project to evaluate how fine-grained FP8 scaling affects both training efficiency and downstream legal reasoning performance. During training, model weights were kept in BF16, while the inputs, weights, and gradient outputs were dynamically quantized to FP8-E4M3 using TorchAO’s rowwise configuration, where each row receives its own scaling factor. This finer granularity enabled higher GPU utilization and reduced training time, achieving ~19.87% speedup over the BF16 baseline on H100 GPUs while maintaining stable convergence. The model was trained on the newmindai/EuroHPC-Legal dataset (multi-domain Q/A format) to improve reasoning quality across various subfields of Turkish law.

Experiment Context

This model was trained with the finer grained resolution for Float8 Rowwise recipe. The recipe sets scaling_granularity to AXISWISE for each of input, weight and gradient output cast configurations and ScalingType to DYNAMIC where each row of a weight matrix gets its own scaling factor, instead of one scaling factor for the entire tensor. and the dtype to float8_e4m3, basically it points that the bites be distributed with 4 exponents and 3 mantissas and the first bit is signal bit.

from torchao.float8 import (
    convert_to_float8_training,
    Float8LinearConfig)
config = Float8LinearConfig.from_recipe_name("rowwise")
model = convert_to_float8_training(model, config=config)

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: 1
• Per-Device Micro Batch Size: 2
• Gradient Accumulation: 4
• GPUs: 16 (via CUDA_VISIBLE_DEVICES=0,1,2,3)
• dtype: bf16 && fp8=false
  • Weights: bfloat16
  • Activations: bfloat16
• 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
• Evaluation: every 5 steps (eval_steps=5, eval_samples=1000)
• 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 (multi 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) 4 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

The training progress has been profiled using pytorch-profiler tool.

• follow the steps to visualize the profiles:
  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="./Llama-3.1-8B-Instruct_rowwise" --port="6006"
  

Usage

Note: the final model has been 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-rw-4nodes"
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}}
}