> This page is for version Nightly (default). > For other versions, use one of these documentation indexes: > - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt > - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt > - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt > - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt > For clean Markdown of any page, append .md to the page URL. > For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt. > For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server. # Beam-1.0-501B-A23B > Fine-tune Beam-1.0-501B-A23B, Reflection AI's 501B sparse MoE language model, with NeMo AutoModel: a one-node LoRA recipe and an 8-node full-parameter recipe. [Beam-1.0-501B-A23B](https://huggingface.co/reflection/Beam-1.0-501B-A23B) is Reflection AI's text-only sparse Mixture-of-Experts (MoE) language model: 501.3B total parameters with 23.4B activated per token. NeMo AutoModel trains it with a native implementation on the shared MoE stack, with expert, pipeline and context parallelism, packed sequences, and a Hugging Face state-dict adapter. > **Info** > > | | | > | ---------------------- | ------------------------------------------------------------------------------- | > | **Task** | Text Generation (MoE, mixed sliding-window and full attention) | > | **Architecture** | `BeamForCausalLM` (`model_type: beam`) | > | **Parameters** | 501B total, 23B activated per token | > | **Validated hardware** | LoRA: 1 node, 8 H200 GPUs, EP8; full-parameter SFT: 8 nodes, 64 H200 GPUs, EP64 | > | **HF Org** | [reflection](https://huggingface.co/reflection) | ## Fine-Tune Beam Follow the [installation instructions](/get-started/installation), then run the LoRA recipe from the repository root on one node of 8xH200 GPUs (141 GB VRAM): ```bash uv run automodel --nproc-per-node=8 examples/llm_finetune/beam/beam_1_0_501b_a23b_tulu3_lora_ep8.yaml ``` The recipe trains rank-32 LoRA adapters on the attention and dense-MLP projections and on the routed experts with EP8, packed 2048-token sequences of the Tulu-3 SFT mixture, and a global batch of 64 packs. For full-parameter SFT, run the 8-node recipe on each of 8 nodes of 8xH200 GPUs, with `NODE_RANK` 0-7 and `MASTER_ADDR` set to the first node: ```bash uv run torchrun --nnodes=8 --nproc-per-node=8 --node-rank=$NODE_RANK --master-addr=$MASTER_ADDR --master-port=29500 \ -m nemo_automodel.cli.app examples/llm_finetune/beam/beam_1_0_501b_a23b_tulu3_ep64_8nodes.yaml ``` It trains all weights with EP64 on the same Tulu-3 data, packed to 4096 tokens, with a global batch of 64 packs. TE FusedAdam keeps fp32 master weights and Adam moments while the model weights and checkpoints stay bf16. Each checkpoint holds about 6 TB, so `checkpoint_dir` must be on a shared filesystem. ## Choose a Workflow | Goal | Start Here | | ------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------- | | Run the primary recipe for this model | Use the [recipe configuration](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/llm_finetune/beam/beam_1_0_501b_a23b_tulu3_lora_ep8.yaml). | | Run full-parameter SFT on 8 nodes | Use the [full-parameter recipe](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/llm_finetune/beam/beam_1_0_501b_a23b_tulu3_ep64_8nodes.yaml). | | Prepare your environment | Follow the [installation instructions](/get-started/installation). | ### Configuration | Setting | Configuration | | -------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------- | | Hardware | 8x H200 (141 GB) | | Strategy | FSDP2; tp_size=1, pp_size=1, cp_size=1, ep_size=8 | | Nodes | 1 | | Features | Checkpointing (enabled=true), Activation checkpointing (activation_checkpointing=true), LoRA, Packed sequences (packed_sequence_size=2048), Transformer Engine | | Advanced | bfloat16, attention=te, linear=torch, experts=torch_mm, dispatcher=torch | ## Model Reference ### Model Architecture | Property | Value | | ---------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------- | | Layers | 52 (3 dense MLP + 49 MoE) | | Hidden size | 5120 | | Attention | GQA, 48 query / 8 KV heads, head dim 128, q/k RMSNorm, sigmoid output gate per head | | Attention layout | 43 sliding-window layers (window 512, RoPE) and 9 full-attention layers (NoPE with scalable-softmax query scaling) | | MoE | 256 routed experts, top-8, no shared expert, expert FFN 2560 (SwiGLU) | | Router | Sigmoid scores with an `e_score_correction_bias` for selection, grouped top-k (8 groups, top 4), normalized weights, scaling factor 2.5, computed in fp32 | | Norms | Sandwich RMSNorm around attention and MLP, depth-scaled residual adds | | Vocabulary | 151,936, untied embeddings and LM head | | Output | Final logit soft-capping at 30.0 | ## Architecture * `BeamStateDictAdapter` converts the per-expert `gate_proj` / `up_proj` / `down_proj` weights of the Hugging Face checkpoint to the grouped expert layout and back, so exports keep the Hugging Face layout. * The router weight and `e_score_correction_bias` stay fp32 through the model cast, checkpoint load, FSDP2 sharding (their own fp32 unit) and export. * Parallelism: EP, FSDP2, PP and CP are supported; TP is not. Context parallelism requires the TE attention backend. Full-attention layers use TE's `p2p` exchange; sliding-window layers use `all_gather`, or `a2a` for packed (THD) input, which needs the CP size to divide the 48 query and 8 KV heads (CP 2, 4 or 8). * Backends: attention `te`, `sdpa` or `flex`, linear `torch` or `te`, experts `torch` or `torch_mm`. `flex` handles only unpadded, unpacked batches: it rejects a padding `attention_mask` and packed (THD) input, so use `te` or `sdpa` for padded or packed data. FP8 uses the existing torchao and TE paths; the grouped experts stay bf16. ## Validated Configuration | Area | Validated Configuration | | ----------------- | ------------------------------------------------------------------------------------------------- | | Training | Full-parameter SFT with FSDP2 and activation checkpointing | | Optimizer | TE FusedAdam with fp32 master weights (bf16 weights plus 16-bit remainders) and fp32 Adam moments | | Distributed setup | EP64 / CP1, TP1 / PP1 | | Workload | `allenai/tulu-3-sft-mixture`, packed 4,096-token sequences | | Batch | Local batch size 1; global batch size 64 | | Scale | 8 nodes / 64 H200 GPUs | | Duration | 100 optimizer steps | The [full-parameter recipe](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/llm_finetune/beam/beam_1_0_501b_a23b_tulu3_ep64_8nodes.yaml) reduced training loss from 0.6622 at step 0 to 0.5137 at step 99, with a mean of 0.5183 over steps 10-99, and the validation loss from 0.4832 after 25 steps to 0.4682 after 100. Peak logged memory was 117.0 GiB per GPU. The run completed without non-finite losses or gradient norms and saved the model and optimizer state after 50 and 100 steps, about 6 TB per checkpoint. Resuming from the checkpoint saved after 50 steps on the same 8 nodes reproduced the next step's loss bit-exactly, stayed within 0.05% of the uninterrupted run over the following 10 steps, and did not raise peak memory. This validates the training mechanics, not model convergence. ## Numerical Validation ### Hugging Face Logit Parity Logits of the native implementation against the Hugging Face reference on 4 held-out Tulu-3 conversations of 4096 tokens (16,384 positions, full vocabulary). Truncated models keep the first N decoder layers plus embeddings, final norm and LM head. | Model | Compute | Mean KL | Logits cosine | Top-1 agreement | | --------------------- | ------- | ----------: | ------------: | --------------: | | First 6 layers | FP32 | 0.000000026 | 0.99999999 | 100% | | First 6 layers | BF16 | 0.0012 | 0.99995 | 90.7% | | Full model, 52 layers | BF16 | 0.0053 | 0.9927 | 98.4% | The FP32 row shows numerical equivalence; the BF16 rows measure rounding noise, which grows with depth and with near-tie expert selection. Loading the full checkpoint maps all 38,366 tensors with no missing or mismatched entries, and a load, sharded save and offline consolidation round trip reproduces all 38,366 tensors bit-exactly at EP1, EP2 and EP8. ### 100-Step LoRA Training The LoRA recipe, [`beam_1_0_501b_a23b_tulu3_lora_ep8.yaml`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/llm_finetune/beam/beam_1_0_501b_a23b_tulu3_lora_ep8.yaml), trained for 100 steps on 1 node / 8 H200 GPUs. | Metric | Result | | ------------------------------------- | --------------: | | Step 0 loss | 0.7253 | | Step 99 loss | 0.4777 | | Mean loss, steps 10-99 | 0.5110 | | Validation loss, after 10 / 100 steps | 0.5272 / 0.4945 | | Peak logged memory | 129.7 GiB | The validation loss on 256 held-out conversations fell at every evaluation. The run completed without non-finite losses or gradient norms. ### Packed CP1 / CP4 / CP8 Training Parity The full model trained for 100 steps with the same Tulu-3 workload, global batch size 64, EP64, TP1 and PP1, using torch AdamW with bf16 optimizer state. The CP1 reference packs documents with the same per-document alignment as the CP runs. | Run | Resources | Peak Logged Memory | Step 99 Loss | | --- | -----------------: | -----------------: | -----------: | | CP1 | 8 nodes / 64 H200s | 74.30 GiB | 0.6219 | | CP4 | 8 nodes / 64 H200s | 92.67 GiB | 0.6213 | | CP8 | 8 nodes / 64 H200s | 91.55 GiB | 0.6219 | Across all 100 matched steps, the CP4 loss curve differed from CP1 by a mean absolute difference of 0.000413 (RMSE 0.000465, maximum 0.0012), and the CP8 curve by 0.000582 (RMSE 0.000880, maximum 0.0022). ## Current Scope * Tensor parallelism is not supported, and the LM head must stay untied from the embeddings. * Packed context parallelism needs CP 2, 4 or 8. ## Related Resources * [LLM fine-tuning guide](/recipes-e2e-examples/sft-peft) * [Mixed-precision training](/development/mixed-precision-training)