From cec7bae44e4d516cabb57e33f439bfa02d1c08d8 Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Tue, 7 Jul 2026 20:44:16 +0000 Subject: [PATCH 1/5] Add GPTModel to HybridModel migration guide Signed-off-by: Philip Petrakian --- docs/index.md | 1 + docs/user-guide/hybrid-model-migration.md | 341 ++++++++++++++++++ docs/user-guide/index.md | 1 + skills/mcore-migrate-gpt-to-hybrid/SKILL.md | 31 ++ .../agents/openai.yaml | 4 + 5 files changed, 378 insertions(+) create mode 100644 docs/user-guide/hybrid-model-migration.md create mode 100644 skills/mcore-migrate-gpt-to-hybrid/SKILL.md create mode 100644 skills/mcore-migrate-gpt-to-hybrid/agents/openai.yaml diff --git a/docs/index.md b/docs/index.md index 11337315588..623f1514614 100644 --- a/docs/index.md +++ b/docs/index.md @@ -50,6 +50,7 @@ get-started/quickstart user-guide/data-preparation user-guide/training-examples user-guide/parallelism-guide +user-guide/hybrid-model-migration ``` ```{toctree} diff --git a/docs/user-guide/hybrid-model-migration.md b/docs/user-guide/hybrid-model-migration.md new file mode 100644 index 00000000000..c24006e06ad --- /dev/null +++ b/docs/user-guide/hybrid-model-migration.md @@ -0,0 +1,341 @@ + + +# Migrate from GPTModel to HybridModel + +This guide describes how to replace a Megatron Core `GPTModel` with a +`HybridModel`, convert an existing distributed checkpoint, and start or resume +training with the converted weights. The conversion stays in Megatron's +distributed-checkpoint format; it does not use Hugging Face as an intermediate +format. + +## 1. What Is HybridModel? + +A standard `GPTModel` decoder layer contains both a self-attention sublayer and +an MLP or MoE sublayer under one layer index. `HybridModel` instead builds an +ordered stack in which every position represents one layer family. The order is +described by `--hybrid-layer-pattern`: + +| Symbol | Layer family | +|--------|--------------| +| `M` | Mamba-2 state-space layer | +| `G` | Gated Delta Network (GDN) layer | +| `*` | Self-attention layer | +| `D` | DeepSeek Sparse Attention (DSA) layer | +| `-` | Dense MLP layer | +| `E` | Mixture-of-Experts (MoE) layer | + +One pattern symbol is one HybridModel layer. Consequently, one GPT transformer +block becomes two HybridModel layers when preserving the original architecture: + +| Source architecture | Equivalent HybridModel pattern | Hybrid layer count | +|---------------------|--------------------------------|--------------------| +| Two dense GPT blocks | `*-*-` | 4 | +| Two all-layer MoE GPT blocks | `*E*E` | 4 | + +For example, source GPT layer 0 is split between Hybrid layers 0 and 1: +its attention parameters move to the first `*`, and its MLP parameters move to +the first `-` or `E`. Source GPT layer 1 maps to the next pair, and so on. + +The pattern can also describe execution layout. A `|` marks a pipeline segment +boundary, and `/` introduces a repeated Multi-Token Prediction (MTP) pattern. +For example, `*-*-|*-*-` places four GPT-equivalent blocks across two pipeline +segments. Separators do not count as layers. + +HybridModel provides the following benefits: + +- Different layer families can be composed in one model without forcing every + decoder block to have the same structure. +- Attention and dense or expert MLP layers can be placed and configured + independently. +- Mamba, GDN, standard or DeepSeek attention, dense MLP, and MoE layers can use + one pattern-driven model interface. +- Patterns that include Mamba can replace some quadratic attention layers with + subquadratic sequence mixing and fixed-size recurrent inference state. +- Pipeline and virtual-pipeline segmentation can be expressed with the model + pattern instead of a separate layer layout. +- The same model abstraction can describe a pure transformer (`*-` repeated), + a pure Mamba model, or a heterogeneous architecture. + +These capabilities do not imply an automatic throughput or quality improvement. +An architecture-preserving `*-` or `*E` migration should be validated for +numerical equivalence, and a pattern that adds another layer family should be +treated as a new architecture and benchmarked independently. + +## 2. How to Convert a Checkpoint + +Use +[`tools/checkpoint/gpt_hybrid_conversion.py`](https://github.com/NVIDIA/Megatron-LM/blob/main/tools/checkpoint/gpt_hybrid_conversion.py) +to convert a `GPTModel` checkpoint directly to `HybridModel` state-dict keys. + +### Choose an architecture-preserving pattern + +For a source checkpoint with *N* GPT layers: + +- Use `*-` repeated *N* times for a dense GPT model. +- Use `*E` repeated *N* times for a GPT model whose MLP in every layer is MoE. + +The converter maps parameters by occurrence, not merely by numeric layer index: + +| Source parameter | Target parameter | +|------------------|------------------| +| Attention from GPT layer *i* | The *i*-th `*` layer | +| MLP or MoE from GPT layer *i* | The *i*-th `-` or `E` layer | +| Embedding and output weights | Copied without changing their model role | +| `decoder.final_layernorm` | Renamed to `decoder.final_norm` | +| No GPT equivalent | Every `M` layer is initialized from scratch | + +Adding `M` symbols is therefore not an architecture-preserving migration. The +converter retains the source attention and MLP weights, but creates fresh +Mamba-2 parameters for the added positions. + +### Check the prerequisites + +The source must use one of these distributed-checkpoint formats: + +- `torch_dist` +- `fsdp_dtensor` + +Legacy `mp_rank_XX/model_optim_rng.pt` checkpoints are not supported. Convert a +legacy checkpoint to `torch_dist` before running this tool. + +Prefer a top-level checkpoint root containing +`latest_checkpointed_iteration.txt` as `--load-dir`. If `--load-dir` points +directly to a directory containing `metadata.json`, the converter writes a flat +target without a tracker file; the standard training entry point expects a +checkpoint root and tracker. + +Run the converter from the repository root in a Megatron environment. A plain +`python` process is sufficient; `torchrun` and a GPU are not required. The tool +gathers full logical tensors on CPU, so the host must have enough memory for the +unsharded source and target model state dicts. Always use a target directory +that is different from the source directory. + +```{warning} +This is a weights conversion, not a resumable full training-state conversion. +Sharded optimizer, RNG, rerun, and Transformer Engine `_extra_state` tensors are +not converted. Some non-tensor entries can remain in `common.pt`, but they do +not constitute a converted optimizer or RNG state. Start the converted model +with a fresh optimizer and RNG state. +``` + +### Run the conversion + +The following example converts a four-layer dense GPT model. Its equivalent +HybridModel has the eight-layer pattern `*-*-*-*-`: + +```bash +python tools/checkpoint/gpt_hybrid_conversion.py \ + --direction gpt-to-hybrid \ + --load-dir /path/to/gpt-checkpoints \ + --save-dir /path/to/hybrid-checkpoints \ + --hybrid-layer-pattern '*-*-*-*-' \ + --reset-iterations +``` + +Always quote the pattern because `*` and `|` have special meaning to a shell. +`--input-format auto` and `--output-format auto` are the defaults: the tool +detects the source backend and writes the same backend. `--reset-iterations` +resets the checkpoint iteration, consumed-sample counters, and cached +`train_iters` and `train_samples`; omit it when the new run must retain that +schedule metadata. + +If the target pattern contains `M`, pass dimensions that match the subsequent +training configuration: + +```bash +python tools/checkpoint/gpt_hybrid_conversion.py \ + --direction gpt-to-hybrid \ + --load-dir /path/to/gpt-checkpoints \ + --save-dir /path/to/hybrid-checkpoints \ + --hybrid-layer-pattern 'M*-M*-M*-M*-' \ + --d-model 4096 \ + --mamba-d-state 128 \ + --mamba2-n-groups 8 \ + --mamba2-head-dim 64 \ + --d-conv 4 \ + --reset-iterations +``` + +The converter currently creates Mamba-2-shaped tensors only. Do not select +Mamba 1 even though the CLI accepts `--mamba-version 1`. In addition, +`2 * d_model` must be divisible by `--mamba2-head-dim`; the converter does not +check this constraint. Keep `--d-conv 4` with the documented default Hybrid +stack, because another convolution width requires a matching custom stack spec. + +The number of `*` positions and the number of `-` or `E` positions must both +equal the source GPT layer count. The pattern validator rejects GDN, DSA, and +mixed dense/MoE layouts. When cached training arguments are present, the tool +also rejects interleaved MoE, experimental or linear attention, heterogeneous +block specifications, Multi-Latent Attention, and MTP checkpoints. That +source-feature validation is incomplete when `common.pt` has no cached `args` +or an older checkpoint lacks a field, so verify those features manually. + +The conversion recognizes standard attention and MLP/MoE state-dict keys only. +Other layer-local tensors are omitted. The documented `hybrid_stack_spec` also +uses Transformer Engine's fused layernorm/linear layout; a local or otherwise +non-TE source layout requires a compatible custom Hybrid stack and key +conversion. Always perform the strict-load check described below. + +Do not append an MTP `/...` suffix during conversion. The converter only maps +the main pattern before the first `/`, so it does not create MTP parameters. + +When the source path is a checkpoint root with +`latest_checkpointed_iteration.txt`, the output contains an iteration directory +and a matching tracker file. The saved full-shape tensors can be resharded by a +later Megatron load for a different tensor, pipeline, expert, or FSDP layout. + +## 3. How to Train a Model + +### Update the training command + +Start with the command that trained the GPT model and make these changes: + +1. Replace `pretrain_gpt.py` with `pretrain_hybrid.py`. +2. Remove `--num-layers` and add the same ordered main-layer symbols used for + conversion. Pipeline `|` separators may be added or moved. The command-line + parser derives `num_layers` from the pattern. +3. Select the HybridModel stack specification with + `--spec megatron.core.models.hybrid.hybrid_layer_specs hybrid_stack_spec`. +4. Load the converted weights with a fresh optimizer and write new training + checkpoints to a separate directory. Set `--ckpt-format` to the converter's + `torch_dist` or `fsdp_dtensor` output format. + +A minimal migration of the model and checkpoint arguments looks like this: + +```diff +- torchrun --nproc_per_node=8 pretrain_gpt.py \ +- --num-layers 4 \ +- --load /path/to/gpt-checkpoints \ +- --save /path/to/gpt-checkpoints ++ torchrun --nproc_per_node=8 pretrain_hybrid.py \ ++ --hybrid-layer-pattern '*-*-*-*-' \ ++ --spec megatron.core.models.hybrid.hybrid_layer_specs hybrid_stack_spec \ ++ --ckpt-format torch_dist \ ++ --pretrained-checkpoint /path/to/hybrid-checkpoints \ ++ --load /path/to/new-training-checkpoints \ ++ --save /path/to/new-training-checkpoints +``` + +Keep the existing architecture, optimizer, precision, data, and basic +TP/DP/EP/CP arguments unless this guide identifies a required change. Review +pattern-driven pipeline layout and GPT-specific dataset features separately. + +```{warning} +`pretrain_hybrid.py` does not select `GPTFIMDataset` when `--fim-data` is set. +A GPT training workflow that uses fill-in-the-middle data needs a custom dataset +path or equivalent Hybrid entry-point support before migration. +``` + +With an empty `--load` directory, `--pretrained-checkpoint` loads the converted +weights with finetuning semantics: iteration starts at zero, and optimizer and +RNG state are not restored. After the job writes a checkpoint to `--load`, later +launches resume the new HybridModel training state normally. Alternatively, use +`--load /path/to/hybrid-checkpoints --finetune` for the first launch, then remove +`--finetune` and resume from the new output directory. + +Do not use `--use-checkpoint-args` for the first converted-checkpoint load. The +converter updates the cached layer count and pattern, but it does not rewrite +every GPT-specific model argument, including the old layer specification and +cached `mamba_state_dim`, `mamba_head_dim`, `mamba_num_groups`, and +`mamba_num_heads`. Provide the HybridModel target arguments explicitly. + +For an architecture-preserving migration, keep hidden size, attention heads, +query groups, FFN size, normalization, positional embeddings, vocabulary, +weight tying, and MoE settings identical to the source GPT run. If the pattern +contains `M`, use the following matching option names. + +| Conversion option | Training option | +|-------------------|-----------------| +| `--d-model` | `--hidden-size` | +| `--mamba-d-state` | `--mamba-state-dim` | +| `--mamba2-n-groups` | `--mamba-num-groups` | +| `--mamba2-head-dim` | `--mamba-head-dim` | + +The converter derives the Mamba head count as +`2 * hidden_size / mamba_head_dim`. An explicit training +`--mamba-num-heads` must equal that value. The default Hybrid stack uses a +convolution width of 4 and has no separate training CLI option for it; keep +converter `--d-conv 4` unless a custom stack spec uses a matching width. + +Set `--position-embedding-type` explicitly to preserve the source architecture. +The direct `GPTModelConfig` and `HybridModelConfig` defaults differ, and the +Hybrid training config does not currently carry GPT's `--use-rope-scaling` +setting or `mrope` path. Models that depend on either need corresponding custom +Hybrid support before their checkpoints can be equivalent. + +New Mamba layers begin with random weights. Expect an initial loss discontinuity +and choose learning-rate warmup and validation criteria as you would for a newly +initialized part of a model. + +### Train from scratch + +To initialize every layer from scratch, use the same `pretrain_hybrid.py`, +`--hybrid-layer-pattern`, and `--spec` arguments, but omit +`--pretrained-checkpoint`. Point `--load` and `--save` at the new run directory +if later launches should resume it. Unlike checkpoint conversion, training from +scratch can use compatible layer families supported by the selected HybridModel +stack specification and can include an MTP suffix such as `M*M*/MM/MM`. +Pattern constraints still apply; for example, standard attention `*` and DSA +`D` cannot be used in the same model. + +### Account for expanded layer indices + +Any setting, mapping, or callback indexed by decoder layer must use HybridModel +indices. For the pattern `*E*E`, source layer 0 attention is Hybrid layer 0 and +its MoE is Hybrid layer 1; source layer 1 attention is Hybrid layer 2 and its +MoE is Hybrid layer 3. Expand attention-only lists with inactive entries for +the intervening MLP or MoE positions. + +For pipeline parallelism, add `|` separators without changing the ordered layer +symbols. For example, the converted pattern `*-*-*-*-` can be trained with two +pipeline segments as `*-*-|*-*-`. The number of pipe-delimited segments must be +divisible by `--pipeline-model-parallel-size`. + +The pattern replaces conventional pipeline layout controls. Remove +`--num-layers-per-virtual-pipeline-stage`, +`--num-virtual-stages-per-pipeline-rank`, `--pipeline-model-parallel-layout`, +`--account-for-embedding-in-pipeline-split`, and +`--account-for-loss-in-pipeline-split`. When the pattern contains `|`, also +remove `--decoder-first-pipeline-num-layers` and +`--decoder-last-pipeline-num-layers`. Express virtual-pipeline segmentation +with additional pipe-delimited segments instead. + +Custom providers and conversion mappings also need to account for these API and +state-dict differences: + +- Build or register `HybridModel` instead of `GPTModel`. +- Supply a `hybrid_stack_spec` instead of a GPT transformer-layer spec. +- Set programmatic `num_layers` to the number of layer symbols in the main + pattern; unlike the CLI path, a custom provider might not derive it. +- Map attention and MLP/MoE parameters to their separate Hybrid layer indices. +- Use `decoder.final_norm` in HybridModel mappings instead of + `decoder.final_layernorm`. +- Expand per-layer settings such as attention-window schedules to the full + HybridModel pattern. +- The declarative `HybridModelBuilder` currently rejects virtual pipeline + parallelism. Pipe-defined virtual stages are supported by the + `pretrain_hybrid.py` CLI builder, but custom builder users must avoid VPP or + use a path that explicitly supports it. + +### Validate before scaling up + +Before starting a long run: + +- Load the converted checkpoint strictly and confirm that no model keys or + tensor shapes are missing or unexpected. +- For a `*-` or `*E` migration, compare logits on a fixed batch against the + source GPT model within the expected precision tolerance. +- Run a few training iterations and inspect the loss, gradient norms, and + parameter counts by layer. +- Save and reload one new checkpoint to confirm that the new optimizer and RNG + state resume correctly. + +Logit equivalence is not expected when the pattern adds freshly initialized +`M` layers; validate that case as an architecture change instead. diff --git a/docs/user-guide/index.md b/docs/user-guide/index.md index 2a7ee2eeab9..02d86440ac0 100644 --- a/docs/user-guide/index.md +++ b/docs/user-guide/index.md @@ -23,4 +23,5 @@ data-preparation training-examples parallelism-guide features/index +hybrid-model-migration ``` diff --git a/skills/mcore-migrate-gpt-to-hybrid/SKILL.md b/skills/mcore-migrate-gpt-to-hybrid/SKILL.md new file mode 100644 index 00000000000..6bd657df44f --- /dev/null +++ b/skills/mcore-migrate-gpt-to-hybrid/SKILL.md @@ -0,0 +1,31 @@ +--- +name: mcore-migrate-gpt-to-hybrid +description: Migrate Megatron Core GPTModel checkpoints, model providers, training commands, and layer mappings to HybridModel. Use for GPTModel-to-HybridModel migration, gpt_hybrid_conversion.py, hybrid layer patterns, converted-checkpoint training, or migration validation. Reads the canonical migration document instead of duplicating it. +--- + +# Migrate GPTModel to HybridModel + +## Canonical source + +Read [`docs/user-guide/hybrid-model-migration.md`](../../docs/user-guide/hybrid-model-migration.md) +completely before forming a migration answer or plan, changing code, converting +a checkpoint, or launching training. + +Treat that document as the sole source of truth for migration behavior, +commands, mappings, prerequisites, limitations, and validation. Do not copy +those details into this skill or add a separate reference file. + +If the document and implementation disagree, report the discrepancy before +continuing. When the task authorizes a correction, update the canonical document +first so future human and skill-driven workflows remain aligned. + +## Workflow + +1. Pull the task artifact first, such as the checkpoint metadata, model + provider, training command, conversion log, or failure output. +2. Read the canonical migration document completely. +3. Follow the relevant document section without inventing unsupported migration + paths or silently changing architecture. +4. Validate the result in proportion to the change, using the repository's + build and testing skills when applicable. +5. Report the outcome and link the canonical document for human readers. diff --git a/skills/mcore-migrate-gpt-to-hybrid/agents/openai.yaml b/skills/mcore-migrate-gpt-to-hybrid/agents/openai.yaml new file mode 100644 index 00000000000..b42df57817e --- /dev/null +++ b/skills/mcore-migrate-gpt-to-hybrid/agents/openai.yaml @@ -0,0 +1,4 @@ +interface: + display_name: "Migrate GPTModel to HybridModel" + short_description: "Migrate GPTModel checkpoints and training" + default_prompt: "Use $mcore-migrate-gpt-to-hybrid to migrate this GPTModel checkpoint and training configuration to HybridModel." From 6203833af0dfe17d892586df41f3c41a28cd7050 Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Wed, 8 Jul 2026 15:24:06 +0000 Subject: [PATCH 2/5] Align GPT-to-Hybrid skill with repository format Signed-off-by: Philip Petrakian --- skills/mcore-migrate-gpt-to-hybrid/SKILL.md | 50 ++++++++++++------- .../agents/openai.yaml | 4 -- 2 files changed, 31 insertions(+), 23 deletions(-) delete mode 100644 skills/mcore-migrate-gpt-to-hybrid/agents/openai.yaml diff --git a/skills/mcore-migrate-gpt-to-hybrid/SKILL.md b/skills/mcore-migrate-gpt-to-hybrid/SKILL.md index 6bd657df44f..486519ecfb9 100644 --- a/skills/mcore-migrate-gpt-to-hybrid/SKILL.md +++ b/skills/mcore-migrate-gpt-to-hybrid/SKILL.md @@ -1,31 +1,43 @@ --- name: mcore-migrate-gpt-to-hybrid -description: Migrate Megatron Core GPTModel checkpoints, model providers, training commands, and layer mappings to HybridModel. Use for GPTModel-to-HybridModel migration, gpt_hybrid_conversion.py, hybrid layer patterns, converted-checkpoint training, or migration validation. Reads the canonical migration document instead of duplicating it. +description: Migration guide for moving Megatron Core GPTModel checkpoints, model providers, training commands, and layer mappings to HybridModel. +license: Apache-2.0 +when_to_use: Migrating or reviewing a GPTModel checkpoint or training workflow for HybridModel; choosing or reviewing a hybrid layer pattern; running gpt_hybrid_conversion.py; loading a converted checkpoint; diagnosing GPT-to-Hybrid migration issues; 'migrate GPTModel to HybridModel', 'convert GPT checkpoint to HybridModel', 'hybrid layer pattern'. +metadata: + author: Philip Petrakian --- -# Migrate GPTModel to HybridModel +# GPTModel to HybridModel Migration -## Canonical source +## Answer-First Migration Guidance -Read [`docs/user-guide/hybrid-model-migration.md`](../../docs/user-guide/hybrid-model-migration.md) -completely before forming a migration answer or plan, changing code, converting -a checkpoint, or launching training. +- The canonical source is + [`docs/user-guide/hybrid-model-migration.md`](../../docs/user-guide/hybrid-model-migration.md). +- Read the canonical document completely before answering, planning, reviewing, + editing, converting, or training. +- Keep migration behavior, commands, mappings, prerequisites, limitations, and + validation in the canonical document only. Do not duplicate them in this + skill. -Treat that document as the sole source of truth for migration behavior, -commands, mappings, prerequisites, limitations, and validation. Do not copy -those details into this skill or add a separate reference file. - -If the document and implementation disagree, report the discrepancy before -continuing. When the task authorizes a correction, update the canonical document -first so future human and skill-driven workflows remain aligned. +--- ## Workflow -1. Pull the task artifact first, such as the checkpoint metadata, model - provider, training command, conversion log, or failure output. +1. Pull the task artifact first: checkpoint metadata, model provider or config, + training command, conversion log, diff, or failure output. 2. Read the canonical migration document completely. -3. Follow the relevant document section without inventing unsupported migration - paths or silently changing architecture. -4. Validate the result in proportion to the change, using the repository's - build and testing skills when applicable. +3. Follow only the relevant document sections. Do not invent an unsupported + migration path or silently change the target architecture. +4. Validate the result proportionately, invoking the relevant repository build + and testing skills when applicable. 5. Report the outcome and link the canonical document for human readers. + +--- + +## Documentation Drift + +If the implementation and migration guide disagree: + +1. Report the discrepancy before continuing. +2. If the task authorizes a correction, update the canonical document first. +3. Do not add a competing migration rule to this skill. diff --git a/skills/mcore-migrate-gpt-to-hybrid/agents/openai.yaml b/skills/mcore-migrate-gpt-to-hybrid/agents/openai.yaml deleted file mode 100644 index b42df57817e..00000000000 --- a/skills/mcore-migrate-gpt-to-hybrid/agents/openai.yaml +++ /dev/null @@ -1,4 +0,0 @@ -interface: - display_name: "Migrate GPTModel to HybridModel" - short_description: "Migrate GPTModel checkpoints and training" - default_prompt: "Use $mcore-migrate-gpt-to-hybrid to migrate this GPTModel checkpoint and training configuration to HybridModel." From e2332c306f3eabcb2955a59d04cf0c111c4b3525 Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Wed, 8 Jul 2026 16:07:00 +0000 Subject: [PATCH 3/5] Update guide Signed-off-by: Philip Petrakian --- docs/user-guide/hybrid-model-migration.md | 33 +---------------------- 1 file changed, 1 insertion(+), 32 deletions(-) diff --git a/docs/user-guide/hybrid-model-migration.md b/docs/user-guide/hybrid-model-migration.md index c24006e06ad..3dca0b5afcd 100644 --- a/docs/user-guide/hybrid-model-migration.md +++ b/docs/user-guide/hybrid-model-migration.md @@ -89,11 +89,6 @@ The converter maps parameters by occurrence, not merely by numeric layer index: | MLP or MoE from GPT layer *i* | The *i*-th `-` or `E` layer | | Embedding and output weights | Copied without changing their model role | | `decoder.final_layernorm` | Renamed to `decoder.final_norm` | -| No GPT equivalent | Every `M` layer is initialized from scratch | - -Adding `M` symbols is therefore not an architecture-preserving migration. The -converter retains the source attention and MLP weights, but creates fresh -Mamba-2 parameters for the added positions. ### Check the prerequisites @@ -102,9 +97,6 @@ The source must use one of these distributed-checkpoint formats: - `torch_dist` - `fsdp_dtensor` -Legacy `mp_rank_XX/model_optim_rng.pt` checkpoints are not supported. Convert a -legacy checkpoint to `torch_dist` before running this tool. - Prefer a top-level checkpoint root containing `latest_checkpointed_iteration.txt` as `--load-dir`. If `--load-dir` points directly to a directory containing `metadata.json`, the converter writes a flat @@ -131,7 +123,7 @@ The following example converts a four-layer dense GPT model. Its equivalent HybridModel has the eight-layer pattern `*-*-*-*-`: ```bash -python tools/checkpoint/gpt_hybrid_conversion.py \ +uv run python tools/checkpoint/gpt_hybrid_conversion.py \ --direction gpt-to-hybrid \ --load-dir /path/to/gpt-checkpoints \ --save-dir /path/to/hybrid-checkpoints \ @@ -146,29 +138,6 @@ resets the checkpoint iteration, consumed-sample counters, and cached `train_iters` and `train_samples`; omit it when the new run must retain that schedule metadata. -If the target pattern contains `M`, pass dimensions that match the subsequent -training configuration: - -```bash -python tools/checkpoint/gpt_hybrid_conversion.py \ - --direction gpt-to-hybrid \ - --load-dir /path/to/gpt-checkpoints \ - --save-dir /path/to/hybrid-checkpoints \ - --hybrid-layer-pattern 'M*-M*-M*-M*-' \ - --d-model 4096 \ - --mamba-d-state 128 \ - --mamba2-n-groups 8 \ - --mamba2-head-dim 64 \ - --d-conv 4 \ - --reset-iterations -``` - -The converter currently creates Mamba-2-shaped tensors only. Do not select -Mamba 1 even though the CLI accepts `--mamba-version 1`. In addition, -`2 * d_model` must be divisible by `--mamba2-head-dim`; the converter does not -check this constraint. Keep `--d-conv 4` with the documented default Hybrid -stack, because another convolution width requires a matching custom stack spec. - The number of `*` positions and the number of `-` or `E` positions must both equal the source GPT layer count. The pattern validator rejects GDN, DSA, and mixed dense/MoE layouts. When cached training arguments are present, the tool From 35aa26d5c489a44073796fcdb7a43acf46c4a2a2 Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Wed, 8 Jul 2026 16:39:09 +0000 Subject: [PATCH 4/5] Further update guide Signed-off-by: Philip Petrakian --- docs/user-guide/hybrid-model-migration.md | 45 +---------------------- 1 file changed, 2 insertions(+), 43 deletions(-) diff --git a/docs/user-guide/hybrid-model-migration.md b/docs/user-guide/hybrid-model-migration.md index 3dca0b5afcd..e47c36aaa67 100644 --- a/docs/user-guide/hybrid-model-migration.md +++ b/docs/user-guide/hybrid-model-migration.md @@ -185,9 +185,7 @@ A minimal migration of the model and checkpoint arguments looks like this: - --save /path/to/gpt-checkpoints + torchrun --nproc_per_node=8 pretrain_hybrid.py \ + --hybrid-layer-pattern '*-*-*-*-' \ -+ --spec megatron.core.models.hybrid.hybrid_layer_specs hybrid_stack_spec \ -+ --ckpt-format torch_dist \ -+ --pretrained-checkpoint /path/to/hybrid-checkpoints \ ++ --pretrained-checkpoint /path/to/hybrid-checkpoints \ # first-time only + --load /path/to/new-training-checkpoints \ + --save /path/to/new-training-checkpoints ``` @@ -205,43 +203,7 @@ path or equivalent Hybrid entry-point support before migration. With an empty `--load` directory, `--pretrained-checkpoint` loads the converted weights with finetuning semantics: iteration starts at zero, and optimizer and RNG state are not restored. After the job writes a checkpoint to `--load`, later -launches resume the new HybridModel training state normally. Alternatively, use -`--load /path/to/hybrid-checkpoints --finetune` for the first launch, then remove -`--finetune` and resume from the new output directory. - -Do not use `--use-checkpoint-args` for the first converted-checkpoint load. The -converter updates the cached layer count and pattern, but it does not rewrite -every GPT-specific model argument, including the old layer specification and -cached `mamba_state_dim`, `mamba_head_dim`, `mamba_num_groups`, and -`mamba_num_heads`. Provide the HybridModel target arguments explicitly. - -For an architecture-preserving migration, keep hidden size, attention heads, -query groups, FFN size, normalization, positional embeddings, vocabulary, -weight tying, and MoE settings identical to the source GPT run. If the pattern -contains `M`, use the following matching option names. - -| Conversion option | Training option | -|-------------------|-----------------| -| `--d-model` | `--hidden-size` | -| `--mamba-d-state` | `--mamba-state-dim` | -| `--mamba2-n-groups` | `--mamba-num-groups` | -| `--mamba2-head-dim` | `--mamba-head-dim` | - -The converter derives the Mamba head count as -`2 * hidden_size / mamba_head_dim`. An explicit training -`--mamba-num-heads` must equal that value. The default Hybrid stack uses a -convolution width of 4 and has no separate training CLI option for it; keep -converter `--d-conv 4` unless a custom stack spec uses a matching width. - -Set `--position-embedding-type` explicitly to preserve the source architecture. -The direct `GPTModelConfig` and `HybridModelConfig` defaults differ, and the -Hybrid training config does not currently carry GPT's `--use-rope-scaling` -setting or `mrope` path. Models that depend on either need corresponding custom -Hybrid support before their checkpoints can be equivalent. - -New Mamba layers begin with random weights. Expect an initial loss discontinuity -and choose learning-rate warmup and validation criteria as you would for a newly -initialized part of a model. +launches resume the new HybridModel training state normally. ### Train from scratch @@ -305,6 +267,3 @@ Before starting a long run: parameter counts by layer. - Save and reload one new checkpoint to confirm that the new optimizer and RNG state resume correctly. - -Logit equivalence is not expected when the pattern adds freshly initialized -`M` layers; validate that case as an architecture change instead. From d68043c7ba67ead877b6483e60f433ae36c3a38f Mon Sep 17 00:00:00 2001 From: Philip Petrakian Date: Mon, 13 Jul 2026 22:09:33 +0000 Subject: [PATCH 5/5] Separate pipeline parallelism guidance Signed-off-by: Philip Petrakian --- docs/user-guide/hybrid-model-migration.md | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/docs/user-guide/hybrid-model-migration.md b/docs/user-guide/hybrid-model-migration.md index e47c36aaa67..445a93e81e2 100644 --- a/docs/user-guide/hybrid-model-migration.md +++ b/docs/user-guide/hybrid-model-migration.md @@ -224,6 +224,8 @@ its MoE is Hybrid layer 1; source layer 1 attention is Hybrid layer 2 and its MoE is Hybrid layer 3. Expand attention-only lists with inactive entries for the intervening MLP or MoE positions. +### Configure pipeline parallelism + For pipeline parallelism, add `|` separators without changing the ordered layer symbols. For example, the converted pattern `*-*-*-*-` can be trained with two pipeline segments as `*-*-|*-*-`. The number of pipe-delimited segments must be @@ -238,6 +240,13 @@ remove `--decoder-first-pipeline-num-layers` and `--decoder-last-pipeline-num-layers`. Express virtual-pipeline segmentation with additional pipe-delimited segments instead. +The declarative `HybridModelBuilder` currently rejects virtual pipeline +parallelism. Pipe-defined virtual stages are supported by the +`pretrain_hybrid.py` CLI builder, but custom builder users must avoid VPP or use +a path that explicitly supports it. + +### Update custom providers and conversion mappings + Custom providers and conversion mappings also need to account for these API and state-dict differences: @@ -250,10 +259,6 @@ state-dict differences: `decoder.final_layernorm`. - Expand per-layer settings such as attention-window schedules to the full HybridModel pattern. -- The declarative `HybridModelBuilder` currently rejects virtual pipeline - parallelism. Pipe-defined virtual stages are supported by the - `pretrain_hybrid.py` CLI builder, but custom builder users must avoid VPP or - use a path that explicitly supports it. ### Validate before scaling up