YAML Metadata Warning: empty or missing yaml metadata in repo card (https://huggingface.co/docs/hub/model-cards#model-card-metadata)

TPU-Optimized LLM Training

This repository contains a highly optimized implementation for training Large Language Models (LLMs) on TPU v4-32 hardware. The code is specifically designed to efficiently train a 600 billion parameter model within a 30-day timeframe.

Features

  • TPU v4-32 Optimizations: Specialized code for TPU v4-32 hardware with efficient parallelism strategies
  • Memory Efficiency: Optimized memory usage with gradient checkpointing and efficient attention mechanisms
  • Performance Monitoring: Comprehensive logging and performance tracking
  • Long Context Support: Support for very long sequences (up to 32K tokens)
  • Enhanced Reasoning: Additional reasoning layers for improved model capabilities

Requirements

See requirements.txt for the full list of dependencies. Key requirements:

jax[tpu]==0.4.20
jaxlib==0.4.20
libtpu-nightly
flax==0.7.5

Usage

To train a model, use the tpu_train.py script:

python tpu_train.py \
  --model_size 600b \
  --train_file /path/to/training/data.jsonl \
  --tokenizer_file /path/to/tokenizer.model \
  --batch_size 32 \
  --gradient_accumulation_steps 8 \
  --learning_rate 1.5e-4 \
  --max_steps 500000 \
  --warmup_steps 5000 \
  --max_seq_length 32768 \
  --output_dir /path/to/output \
  --parallelism_type tensor \
  --tensor_parallel_size 8 \
  --use_flash_attention \
  --use_gradient_checkpointing \
  --use_rope_scaling \
  --use_reasoning_layer

Architecture

The implementation includes:

  • Optimized Flash Attention: Blocked implementation for efficient memory usage
  • Tensor Parallelism: Efficient parameter sharding across TPU devices
  • Data Parallelism: Optimized data loading and processing
  • Mixed Precision Training: BFloat16 support for TPU
  • Gradient Checkpointing: Memory-efficient backpropagation

Performance

On TPU v4-32 hardware, this implementation achieves:

  • Efficient training of 600B parameter models
  • Support for sequence lengths up to 32K tokens
  • Memory-efficient operation with gradient checkpointing
  • Optimized communication patterns for TPU pods

License

MIT

Downloads last month

-

Downloads are not tracked for this model. How to track
Inference Providers NEW
This model isn't deployed by any Inference Provider. ๐Ÿ™‹ Ask for provider support