motif3-training-example

Training framework for Motif-3, a Mixture-of-Experts transformer, built on torchtitan. Train-only; targets NVIDIA B200 (SM100). The example configs run 4 nodes × 8 GPUs, so you need at least 4 B200 nodes.

For the model details, see the Motif-3 technical report.

Model

Motif-3 MoE (~13B active / ~314B total parameters). Key components:

• Attention — GDLA / Differential Attention v2 with MLA-style Q/KV LoRA, MHC (manifold-constrained hyper-connections), and interleaved sliding-window + full attention.
• MoE — sigmoid top-8 routing with a shared expert, PolyNorm, and DeepSeek-V3 aux-loss-free load balancing; DeepEP (HybridEP) expert parallelism.
• Training — Multi-Token Prediction (MTP), the Muon optimizer with QK-clip, MXFP8 quantization (DeepGEMM), FlashAttention, and ring / Ulysses sequence parallelism.

Two model flavors, selected via [model] flavor in the TOML:

Layout

• llm_training/ — training code; llm_training.train is the entrypoint, motif/ holds the model + configs, layers/ the attention/MoE/MHC kernels, distributed/ and quantization/ the parallelism and MXFP8 helpers.
• llm_training/motif/train_configs/*.toml — training configs.
• docker/cuda/llm-training.Dockerfile — the B200 training image.
• activations/, optimizer/ — vendored CUDA / Triton kernels, imported as the activation and optimizer packages.
• patches/ — FSDP2 monkey-patches (reduce-scatter overlap, fp32 all-reduce, …).

Image

Jobs run on the prebuilt image ghcr.io/motiftechnologies/llm-training:v0.1.12. It bundles NVSHMEM, DeepEP (HybridEP), DeepGEMM, FlashAttention, TransformerEngine, the training dependencies, and the Liger-Kernel fork — everything except the vendored activation package, which each job builds at startup (see Launch). To build the image yourself:

docker build -f docker/cuda/llm-training.Dockerfile -t <tag> .
# or, with the bake file:
docker buildx bake -f docker-bake.hcl

Configure

Copy an example config and edit it — everything is set in the TOML, so the launch command needs no CLI overrides:

• pretrain: llm_training/motif/train_configs/example_c4_pretrain.toml (motif3_seq)
• SFT: llm_training/motif/train_configs/example_c4_sft.toml (motif3_seq_256k_mtp1_sft)

Both are laid out for 4 nodes × 8 GPU (32 ranks).

Set [model] flavor, the [parallelism] degrees, and [training] batch/seq/steps. Data loading is HuggingFace-only:

[training]
dataset = "huggingface"
dataset_path = "/path/to/dataset"   # local *.jsonl or a datasets.save_to_disk arrow dir

[huggingface_dataset]
split = "train"

Samples are read in chat form ({"conversations": [{"role": ..., "content": ...}]}) and rendered with the tokenizer's chat template. Only the last assistant turn is trained on, so pretraining-style raw text must sit in an assistant turn.

Launch

Submit a Kubeflow TrainJob on at least 4 B200 nodes (8 GPUs each). Before training, each node builds the activation package — the one package the image does not bake. Build it from a node-local copy: building ./activations in place drops build/ + egg-info into the checkout, and when the checkout is on shared storage, concurrent node builds race and fail with "Failed to build installable wheels for activation":

ACT_TMP=$(mktemp -d)
cp -r ./activations "$ACT_TMP/activations"
rm -rf "$ACT_TMP/activations/build" "$ACT_TMP/activations/activation.egg-info"
pip install --no-build-isolation "$ACT_TMP/activations"
rm -rf "$ACT_TMP"

Then each node runs:

CONFIG_FILE=<config.toml> bash run_train_kubeflow.sh

run_train_kubeflow.sh is a thin torchrun -m llm_training.train --job.config_file $CONFIG_FILE wrapper (node / proc counts come from the TrainJob spec). Any --section.key=value arguments appended to it override the TOML.

Weights & Biases

Enable wandb in the config and pass credentials via the environment (use a k8s Secret for the key):

[metrics]
enable_wandb = true

WANDB_PROJECT=... WANDB_TEAM=... WANDB_RUN_NAME=... WANDB_API_KEY=...

License

MIT — see LICENSE. Vendored third-party code keeps its original license; THIRD_PARTY_NOTICES.md carries the full texts: torchtitan / PyTorch (BSD-3-Clause), ring-flash-attention (MIT), matmul_transpose_triton (MIT), and the vLLM / kernel-builder-derived activations/ kernels (Apache-2.0).