From cf41952eb5ba7afc76e10ec31e22b9d44ab17360 Mon Sep 17 00:00:00 2001 From: Jinze Xue Date: Mon, 11 May 2026 12:03:44 -0700 Subject: [PATCH 1/4] tokenizer: add --null-tokenizer-vocab-includes-eod When set, NullTokenizer treats --vocab-size N as the total vocab including eod, so eod_id=N-1 and tokenizer.vocab_size=N. Default behavior is unchanged (eod_id=N, tokenizer.vocab_size=N+1). Matches Megatron-Bridge's NullTokenizer convention for mock-data benchmarks. Co-Authored-By: Claude Opus 4.7 (1M context) --- megatron/core/tokenizers/utils/build_tokenizer.py | 5 ++++- megatron/training/arguments.py | 8 ++++++++ 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/megatron/core/tokenizers/utils/build_tokenizer.py b/megatron/core/tokenizers/utils/build_tokenizer.py index bf02451ae6c..2be563dc714 100644 --- a/megatron/core/tokenizers/utils/build_tokenizer.py +++ b/megatron/core/tokenizers/utils/build_tokenizer.py @@ -73,7 +73,10 @@ def build_tokenizer(args, **kwargs): ) metadata = {'library': tokenizer_library} if args.vocab_size: - kwargs['vocab_size'] = args.vocab_size + if getattr(args, 'null_tokenizer_vocab_includes_eod', False): + kwargs['vocab_size'] = args.vocab_size - 1 + else: + kwargs['vocab_size'] = args.vocab_size tokenizer = MegatronTokenizer.from_pretrained(metadata_path=metadata, **kwargs) # Add vocab size (if not already set from a checkpoint). diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 6a108a0d6d0..4e9dcbbb7e6 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -3974,6 +3974,14 @@ def _add_tokenizer_args(parser): group.add_argument( '--tokenizer-model', type=str, default=None, help='Sentencepiece tokenizer model.' ) + group.add_argument( + '--null-tokenizer-vocab-includes-eod', + action='store_true', + default=False, + help='Treat --vocab-size as the total vocab including eod. ' + 'When set, NullTokenizer eod_id=N-1 (instead of default N). ' + 'For compatibility with Megatron-Bridge convention.', + ) group.add_argument( '--tokenizer-metadata', type=str, From f69baebd148856496acd99c5e1eac83ddc945408 Mon Sep 17 00:00:00 2001 From: jinzex Date: Wed, 1 Apr 2026 19:27:19 -0700 Subject: [PATCH 2/4] Add TP-invariant numerics for Megatron-Core MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bitwise-identical forward, backward, and end-to-end training across TP=1, 2, 4, 8 on Megatron-Core TransformerBlocks. Gated by NVTE_TP_INVARIANT_MODE environment variable. Components: - TP-invariant GEMM (all-gather sharded weight, full-K GEMM) for both column and row parallel linear in the TE patches under examples/tp-numerics/patches/. - Gradient clipping pow2-rounding to absorb 1-ULP cross-TP norm jitter. - RMSNorm dgamma all-gather with rank-0-only reduction. - Batch-invariant Triton kernels (BIK) for M-invariant matmul. - Cross-entropy all-gather over exp_logits. Validation: - tests/unit_tests/transformer/test_tp_invariant.py — TP=1≡2≡4 bitwise on a small TransformerBlock (fp32+bf16). - examples/tp-numerics/submit_qwen3_{0.6b,8b,moe_toy}_tp_invariant.sh — end-to-end raw-MLM training scripts; Qwen3-0.6B (TP=1) and Qwen3-8B (TP=4) bitwise across 100 iters and across B300 ≡ H100 hardware. Co-Authored-By: Claude Opus 4.7 (1M context) --- examples/tp-numerics/.gitignore | 12 + examples/tp-numerics/README.md | 167 ++ .../assets/b300_vs_h100_tp1_invariant.png | Bin 0 -> 126615 bytes .../assets/invariant_vs_baseline_tp4.png | Bin 0 -> 80005 bytes .../assets/tp4_vs_tp8_invariant.png | Bin 0 -> 65773 bytes examples/tp-numerics/diff_dumps.py | 45 + examples/tp-numerics/patches/README.md | 35 + .../dot_product_attention/backends.py | 1948 +++++++++++++++++ .../dot_product_attention.py | 1560 +++++++++++++ .../pytorch/module/layernorm_linear.py | 1913 ++++++++++++++++ .../pytorch/module/linear.py | 1827 ++++++++++++++++ examples/tp-numerics/plot_loss.py | 68 + .../submit_qwen3_0.6b_tp_invariant.sh | 182 ++ .../submit_qwen3_8b_tp_invariant.sh | 181 ++ .../submit_qwen3_moe_toy_tp_invariant.sh | 192 ++ .../core/extensions/transformer_engine.py | 11 +- megatron/core/optimizer/clip_grads.py | 77 +- .../core/tensor_parallel/cross_entropy.py | 23 +- megatron/core/tensor_parallel/layers.py | 36 +- .../custom_layers/batch_invariant_kernels.py | 158 +- .../core/transformer/transformer_config.py | 6 +- megatron/training/training.py | 6 +- .../test_te_layers_batch_invariant.py | 155 ++ .../transformer/test_tp_invariant.py | 211 ++ 24 files changed, 8779 insertions(+), 34 deletions(-) create mode 100644 examples/tp-numerics/.gitignore create mode 100644 examples/tp-numerics/README.md create mode 100644 examples/tp-numerics/assets/b300_vs_h100_tp1_invariant.png create mode 100644 examples/tp-numerics/assets/invariant_vs_baseline_tp4.png create mode 100644 examples/tp-numerics/assets/tp4_vs_tp8_invariant.png create mode 100644 examples/tp-numerics/diff_dumps.py create mode 100644 examples/tp-numerics/patches/README.md create mode 100644 examples/tp-numerics/patches/v2.9.0/transformer_engine/pytorch/attention/dot_product_attention/backends.py create mode 100644 examples/tp-numerics/patches/v2.9.0/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py create mode 100644 examples/tp-numerics/patches/v2.9.0/transformer_engine/pytorch/module/layernorm_linear.py create mode 100644 examples/tp-numerics/patches/v2.9.0/transformer_engine/pytorch/module/linear.py create mode 100644 examples/tp-numerics/plot_loss.py create mode 100644 examples/tp-numerics/submit_qwen3_0.6b_tp_invariant.sh create mode 100644 examples/tp-numerics/submit_qwen3_8b_tp_invariant.sh create mode 100644 examples/tp-numerics/submit_qwen3_moe_toy_tp_invariant.sh create mode 100644 tests/unit_tests/transformer/test_tp_invariant.py diff --git a/examples/tp-numerics/.gitignore b/examples/tp-numerics/.gitignore new file mode 100644 index 00000000000..540ffde29e7 --- /dev/null +++ b/examples/tp-numerics/.gitignore @@ -0,0 +1,12 @@ +# Result logs (kept local, not in upstream PR) +results/ +output/ +*.log + +# Bridge-based legacy scripts (replaced by submit_qwen3_*_tp_invariant.sh) +validate_e2e_*.py +submit_e2e_*.sh +test_tp_numerics.py + +# Local-only notes (never push) +*.local.md diff --git a/examples/tp-numerics/README.md b/examples/tp-numerics/README.md new file mode 100644 index 00000000000..04f36aa9dcd --- /dev/null +++ b/examples/tp-numerics/README.md @@ -0,0 +1,167 @@ +# TP-Invariant Numerics: Bitwise Identical Training Across TP Degrees + +Bitwise identical forward, backward, and E2E training for Megatron-Core TransformerBlocks +regardless of Tensor Parallelism (TP) degree — TP=1, 2, 4, 8 produce the same result. + +**Source branches** (view diffs): +- [Megatron-LM](https://github.com/jinzex/Megatron-LM/tree/jinzex/tp-invariant-numerics) — MCore changes (clip_grads, batch_invariant_kernels, etc.) +- [TransformerEngine](https://github.com/jinzex/TransformerEngine/tree/jinzex/tp-invariant-numerics) — TE changes (layernorm_linear, linear, etc.) + +## Status + +| Model | Unit Test (single fwd+bwd) | E2E Training | +|-------|-----------|--------------| +| **Dense** | TP=1/2/4/8 bitwise identical | TP=1/2/4 bitwise identical loss+grad_norm, **100 iters** | +| **MoE** | TP=1/2/4/8 bitwise identical (BIK) | Pending | + +Config: `BIK=1 NVTE_TP_INVARIANT_MODE=1` | TE 2.9 | Qwen3-0.6B / 8B | H100, B300 + +**Cross-architecture:** B300 ≡ H100 bitwise (Qwen3-0.6B TP=1, Qwen3-8B TP=4, 100 iters). See [Cross-Architecture Validation](#cross-architecture-validation-b300-vs-h100). + +## Baseline vs TP-Invariant (Qwen3-0.6B) + +Without TP-invariant mode, loss diverges from iteration 1: + +| Iter | Baseline TP=1 | Baseline TP=2 (diverges) | Baseline TP=4 (diverges) | TP-Inv TP=1/2/4 (all identical) | +|------|--------------|--------------------------|--------------------------|-------------------------------| +| 1 | 1.213320E+01 | 1.213298E+01 | 1.213321E+01 | **1.213320E+01** | +| 2 | 1.187939E+01 | 1.188345E+01 | 1.188395E+01 | **1.187872E+01** | +| 5 | 1.252811E+01 | 1.253201E+01 | 1.253240E+01 | **1.257282E+01** | +| 10 | 1.165137E+01 | 1.164608E+01 | 1.164373E+01 | **1.143649E+01** | +| 100 | 8.277067E+00 | - | - | **8.277067E+00** | + +## Components & Patches + +All fixes are required together. TE patches must be copied into the container at runtime +(TE is a site-package). MCore changes are committed directly to this branch. + +TE-side patches live in the companion TE branch +[`jinzex/tp-invariant-numerics`](https://github.com/jinzex/TransformerEngine/tree/jinzex/tp-invariant-numerics); +MCore changes are committed in this branch. + +| Component | Fix | Location | +|-----------|-----|----------| +| **TP-Invariant GEMM** (fwd+bwd) | All-gather sharded weight, full-K GEMM | TE `transformer_engine/pytorch/module/{layernorm_linear,linear}.py` | +| **Gated deinterleave** (bwd) | Reorder via `partition_stride` after all-gather | TE `transformer_engine/pytorch/module/layernorm_linear.py` | +| **Cross-entropy** (fwd) | All-gather exp_logits, local sum | `megatron/core/tensor_parallel/cross_entropy.py` | +| **Output projection** (bwd) | All-gather weight+grad, full dgrad GEMM | `megatron/core/tensor_parallel/layers.py` | +| **Gradient clipping** | Float64 norm + pow2 clip_coeff rounding | `megatron/core/optimizer/clip_grads.py` | +| **RMSNorm dgamma** (bwd) | All-gather tokens + rank-0-only reduction | `megatron/core/transformer/custom_layers/batch_invariant_kernels.py` | +| **BIK** (fwd+bwd) | M-invariant Triton matmul_persistent | `megatron/core/transformer/custom_layers/batch_invariant_kernels.py` | + +### Gradient Clipping Fix + +`multi_tensor_l2norm` (Apex CUDA) reduces over different-shaped TP shards → different +FP32 partial sums → different clip_factor. Fix: float64 norm computation + pow2 rounding +of clip_coeff. Mismatch rate: 40.3% (float32) → 0.000% (pow2). + +## Quick Start + +**Prerequisites:** TE 2.9, 1+ GPU (1 for 0.6B-TP1, 4 for 8B-TP4, 8 for MoE-toy), this Megatron-LM branch on PYTHONPATH. **No Bridge dependency** — dense scripts use raw MLM `pretrain_gpt.py` directly. + +### TE patches + +Submit scripts install the matching TE patches into the container at launch +(squashfs containers are read-only per-srun, so patches reapply each run). +Version-pinned snapshots live under `patches/v/` mirroring the TE +source tree — see [`patches/README.md`](patches/README.md) for the supported +versions and how to refresh from the companion +[TE branch](https://github.com/jinzex/TransformerEngine/tree/jinzex/tp-invariant-numerics). +No matching version → submit script fails fast. + +### Run tests + +```bash +cd $PROJ # this Megatron-LM branch + +# Unit test: TP=1 ≡ TP=2 ≡ TP=4 bitwise on a small TransformerBlock +NVTE_TP_INVARIANT_MODE=1 \ +torchrun --nproc_per_node=4 -m pytest \ + tests/unit_tests/transformer/test_tp_invariant.py -v -s + +# E2E: Qwen3-0.6B (100 iters, TP=1; 1 GPU) +bash examples/tp-numerics/submit_qwen3_0.6b_tp_invariant.sh + +# E2E: Qwen3-8B (100 iters, TP=4; 4 GPUs) +bash examples/tp-numerics/submit_qwen3_8b_tp_invariant.sh + +# Smoke test: toy MoE (Qwen3-30B-A3B 4L 8E top-2, 10 iters, 1 GPU) +bash examples/tp-numerics/submit_qwen3_moe_toy_tp_invariant.sh + +# Override TP / iters via env: +# TP_SIZE=2 TRAIN_ITERS=10 bash submit_qwen3_0.6b_tp_invariant.sh +# Each script also runs as `sbatch