nanotron

Nanotron's pipeline parallelism is built on four core components:

• PipelineBlock: Contains model computation split across devices. A block's .rank specifies its assigned pipeline parallel rank.
• PipelineEngine: Orchestrates forward/backward passes. Subclasses override train_batch_iter and validate_batch_iter to implement different schedules like 1F1B or GPipe.
• PipelineBatchState: Manages P2P operations and orchestrates communication across microbatches.
• TensorPointer: A placeholder for tensors produced on different devices. It enables lazy communication, allowing pipeline stages to request tensors on-demand rather than communicating all activations upfront.

Nanotron supports three types of linear layers for tensor parallelism:

1. Regular Column Parallel: Each rank computes only its portion of the output matrix corresponding to its weight shard. It does not compute the full output.
2. Asynchronous Column Parallel: Each rank computes the entire output matrix locally. It kicks off an asynchronous all-gather on the input tensor at the start. While waiting, it computes the portion of the output using its local weight shard. Once the all-gather completes, it computes the remaining portions. This trades higher FLOPs for reduced communication, making it ideal for communication-bound models.
3. Tied Linear: The entire weight matrix is replicated across all ranks rather than being sharded. This is used when weights are shared (e.g., tying embedding and LM head weights).

To run the Mamba modeling example with Nanotron, install the required dependencies and execute the training script provided in the example directory.

1. Install dependencies: pip install -r requirements.txt
2. Run training: ./examples/mamba/train_mamba.sh

To run the LlaMoE example, follow these steps to generate a configuration, install dependencies, and execute the training script using torchrun.

Note: Setting CUDA_DEVICE_MAX_CONNECTIONS=1 is important for certain distributed operations.

# Generate a config file
python examples/moe/config_llamoe.py

# Install megablocks
pip install megablocks

# Run training
export CUDA_DEVICE_MAX_CONNECTIONS=1
torchrun --nproc_per_node=4 examples/moe/train_moe.py --config-file examples/moe/config_llamoe.yaml

Nanotron uses a custom initialization context manager to override PyTorch's default behavior (which often initializes on CPU in FP32). This allows models to be initialized directly on the target GPU in the target precision (e.g., FP16/BF16), saving memory and time.

Mechanism: The context manager overrides nn.Module.register_parameter() and nn.Module.register_buffer() to ensure tensors are created on the correct device and with the correct dtype from the start.

To start training on a single node, use torchrun. You must provide the path to your YAML configuration file via the --config-file flag.

It is recommended to set CUDA_DEVICE_MAX_CONNECTIONS=1 for certain distributed operations. Ensure --nproc_per_node matches your total number of GPUs and aligns with your parallelism settings (dp, tp, pp).

CUDA_DEVICE_MAX_CONNECTIONS=1 torchrun --nproc_per_node=8 run_train.py --config-file examples/config_tiny_llama.yaml

To develop on a remote GPU instance using VS Code:

1. Install the Remote: SSH extension in VS Code.
2. Configure your SSH connection to use your private key. When copying the SSH command from Vast.ai, append -i <path_to_private_key>. Example: ssh -p 50095 root@154.20.254.95 -L 8080:localhost:8080 -i ~/.ssh/id_nanotron
3. Verify your SSH config file (accessible via the config icon in VS Code) includes the identity file path.
4. Once connected, you may need to create a new SSH key specifically for the GPU instance environment:

   ssh-keygen -t rsa
   eval "$(ssh-agent -s)"
   ssh-add

To install Nanotron, create a Python 3.11 virtual environment using uv, install PyTorch with CUDA 12.4 support, and install the core package in editable mode. To run example scripts and use fused kernels, additional dependencies like datasets, transformers, triton, and flash-attn are required.

Note for Hugging Face cluster users: Add export UV_LINK_MODE=copy to your .bashrc to suppress cache warnings from uv.

To use your own model architecture instead of the default LlaMoE, follow these steps:

1. Update the LlaMoEConfig class in config_llamoe.py to match your model's configuration parameters.
2. Update the LlaMoEForTraining class in modeling_llamoe.py to match your model's architecture.
3. Pass your custom model class and config class to the DistributedTrainer in train_moe.py.

trainer = DistributedTrainer(config_file, model_class=LlaMoEForTraining, model_config_class=LlaMoEConfig)

Before running the DoReMi experiment, install the necessary dependencies using the provided requirements file located in the examples/doremi/ directory.

pip install -r examples/doremi/requirements.txt