[feat] FP8 training

[tscholak]

🧐 Problem Description

FP8 training can significantly improve training throughput by reducing memory requirements and improving computational efficiency. However, challenges remain in integrating FP8 across all components of the training stack, such as FlashAttention, while maintaining model performance and stability. For example, the FP8 implementation in FlashAttention 3 raises questions about quantization accuracy and latency between operations.

💡 Proposed Solution

Introduce FP8 training in Fast-LLM with the following steps:

1. Matmul Operations: Implement FP8-compatible matrix multiplication using Transformer Engine or Triton kernels.
2. FlashAttention Compatibility:
   • Evaluate the feasibility of integrating FlashAttention 3, which supports FP8.
   • Address concerns about quantization accuracy and potential pipeline latency.
   • Explore alternatives like FlexAttention if FlashAttention 3 proves unsuitable.
3. Validation:
   • Compare loss curves and convergence behavior between FP8 and BF16 to ensure model stability and performance.

🔄 Alternatives Considered

1. Continue with BF16/FP16 Precision:
   • Pros: Stable and widely supported.
   • Cons: Misses throughput and memory benefits of FP8.
2. Partial FP8 Integration:
   • Pros: Focuses on high-impact components like matrix multiplications.
   • Cons: Could complicate debugging due to mixed precision across components.
3. FlexAttention:
   • Pros: Potentially better support for FP8.
   • Cons: May require more development and testing for integration. May be slower than FlashAttention 3.

📈 Potential Benefits

1. Increased Throughput: FP8 training has demonstrated up to 1.5x speedup compared to BF16 on NVIDIA H100 GPUs.
2. Enhanced Efficiency: Reducing memory overhead enables larger batch sizes and improves GPU utilization.
3. Scalability: Lower precision facilitates training on larger clusters without hitting memory limits.
4. Comparable Convergence: Loss curves for FP8 have been shown to closely track those for BF16, ensuring minimal compromise in training stability.

📝 Additional Context

1. Key Findings from Databricks:
   • FP8 training achieves ~1.5x throughput improvement.
   • Integrated with PyTorch FSDP and activation checkpointing for memory efficiency.
   • Loss curves for FP8 closely match BF16 during training.
   • Databricks Blog: Turbocharged Training with FP8.
2. FlashAttention 3 Insights:
   • FlashAttention 3 supports FP8 but has open questions about latency and accuracy.
   • FlashAttention 3 Issue: FP8 Concerns.
3. Known Challenges:
   • Managing FP8 weight storage and quantization scaling.
   • Addressing pipeline latency issues in FP8 operations.
   • Ensuring alignment between different FP8-compatible components.

[tscholak]

Looks like we can concentrate on linear layers and make it work by broadcasting the quantization scale: https://github.com/pytorch/torchtitan/blob/main/docs/float8.md

[tscholak]

We should use Deepseek's FP8 gemm kernels: https://github.com/deepseek-ai/DeepGEMM

[tscholak]

from the Nemotron-H paper:

Nemotron-H-56B-Base is the first Nemotron model to be fully pre-trained using a FP8-based recipe. Our recipe, called per-tensor current scaling, is coarse grained and involves quantizing entire tensors using a single scaling factor. This factor is chosen to preserve the largest value in the given tensor; any value too small to fit in the desired FP8 format is flushed to zero. In addition, we find it important for training stability and convergence to leave the first and last 4 GEMMs of the model in BF16. In smaller ablations with 8B-parameter models on token horizons up to 15 trillion tokens, we show that FP8 per-tensor scaling can reach equal or better downstream task accuracy compared to BF16.

[weiyx16]

Do you have any idea on FP8 Attention? I can't find any public successful trials about this target.