Conversation
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
Comment on lines
+53
to
+57
| from megatron.bridge.training.mixed_precision import ( | ||
| MixedPrecisionConfig, | ||
| bf16_mixed, | ||
| get_mixed_precision_config, | ||
| ) |
| try: | ||
| with open(out, "w") as fh: | ||
| fh.write(text + "\n") | ||
| except OSError: |
| from __future__ import annotations | ||
|
|
||
| import os | ||
| from typing import Any, Optional |
![pr-pundit-code-review[bot]](https://avatars.githubusercontent.com/u/16860860?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review: AMD-AGI/Primus
Summary
No violations found, 24 files reviewed.
Rules checked: 7 (7 passed, 0 failed).
✅ Passed
All 7 rules satisfied. No architecture concerns found.
Posted by PR Pundit — AI-powered code review grounded in this repo's merge history.
There was a problem hiding this comment.
Code Review: AMD-AGI/Primus
Summary
No violations found, 24 files reviewed.
Rules checked: 7 (7 passed, 0 failed).
✅ Passed
All 7 rules satisfied. No architecture concerns found.
Posted by PR Pundit — AI-powered code review grounded in this repo's merge history.
Implement Megatron-Bridge recipe support for Gemma 4 models to enable SFT/LoRA post-training and pretraining workflows. Implementation: - Recipe module at primus/backends/megatron_bridge/recipes/gemma/ - gemma4.py with pretrain and finetune flavors for both models - Follows upstream Gemma 2 recipe pattern - Model configs: gemma4_26b.yaml and gemma4_31b.yaml - Example configs for SFT and pretraining on MI300X Architecture support: - 26B MoE: 30 layers, 2816 hidden, 128 experts (top-8) - Default: TP=1, EP=8, PP=1 - 31B Dense: 60 layers, 5376 hidden - Default: TP=2, EP=1, PP=1 with sequence parallelism Key features handled by AutoBridge from HF: - Sliding window attention (5 local + 1 global) - GeGLU activation - Dual RoPE timescales - Logit soft capping - MoE routing (26B) Bug fixes applied: - Fix DistributedDataParallelConfig import (megatron.bridge.training.config) - Fix duplicate hf_path parameters in finetune functions - Pop hf_path from user_kwargs to prevent parameter collision
The Gemma 4 recipe added in a7188bd calls AutoBridge.from_hf_pretrained, but the pinned Megatron-Bridge (9577b1280, Jan 2026) has no Gemma 4 bridge registered, so dispatch fails at runtime. Nothing caught this statically: the recipe imports only version-stable public API, and it was patterned on Gemma 2, which does exist at that pin. - Bump third_party/Megatron-Bridge to v0.6.0 (51885cf13). This adds gemma4_bridge.py, gemma4_provider.py and the VL variants, and moves the bundled megatron-core to 0.19.0. v0.5.0 is the first release containing Gemma 4; v0.6.0 additionally carries the GEMMA4_CONVERSION_MODE=text fix for Gemma 4 MoE checkpoints, which the 26B needs. - Pin transformers to 5.12.1 in both bridge hook requirement files. Gemma 4 does not exist anywhere on the 4.x line, and AutoBridge dispatches on the HF architecture class, so a 5.x floor is required rather than preferred. - Force GEMMA4_CONVERSION_MODE=text around both AutoBridge calls. The published Gemma 4 checkpoints are Gemma4ForConditionalGeneration, so the default "auto" dispatch selects Gemma4VLBridge and builds a Gemma4VLModel with vision and audio towers instead of the language model. Mirrors megatron.bridge.recipes.gemma.h100.gemma4. Still requires container validation: a mock-data smoke run, ROCm transformer-engine against megatron-core 0.19.0, and a revalidation pass over the other bridge models that share the transformers pin.
…v0.6.1 Bridge v0.6.1 deleted megatron.bridge.recipes.utils.finetune_utils and moved default_peft_config / default_squad_config into dataset_utils, so importing the Gemma 4 recipe module raised ModuleNotFoundError and no Gemma 4 flavor could be loaded at all. The failure was hard to place because _resolve_recipe() catches ImportError and continues to the next candidate module, so the real cause surfaced only as "Function 'gemma4_26b_pretrain_config' not found".
Gemma4VLBridge.provider_bridge consults GEMMA4_CONVERSION_MODE only on the dense branch. Both published checkpoints declare Gemma4ForConditionalGeneration, so the 26B-A4B (enable_moe_block=true) always fell through to Gemma4VLModelProvider and built vision + audio towers even when the recipe asked for text-only pretraining. Training then died in the loss: the VL model returns the LLaVA-style (loss, new_loss_mask) tuple, but gpt_step.forward_step never passes loss_mask into the model, so masked_next_token_loss dereferenced None. Add the missing MoE branch as a runtime patch, reusing Bridge's own _build_moe_provider so the HF-to-provider field mapping is not duplicated. The upstream fix is two lines in gemma4_vl_bridge.py; patching here avoids writing to the third_party submodule, matching the approach already used for MLPerf. Also add an opt-in Transformer Engine core attention for the dense path, enabled with PRIMUS_GEMMA4_DENSE_ATTENTION_BACKEND=te. Bridge pins dense Gemma 4 to LocalSpecProvider, so it runs plain DotProductAttention with no flash attention and keeps the full score matrix for backward; TE frees enough activation memory to raise micro_batch_size. Only core_attention is swapped, since a full TESpecProvider would replace linear_qkv with a fused-layernorm variant that collides with the dense layer's separate input_layernorm.
Both example configs ran out of memory as written. 31B dense used TP=2, which cannot fit; TP=4 leaves almost no headroom, so the smallest workable degree is TP=8. 26B MoE used TP=1/EP=8 at seq 8192. Expert parallelism shards only the experts, so every rank still holds a full copy of the non-expert weights. TP=2/EP=4 fits; TP=1/EP=8 remains fine at shorter sequence lengths. Enable sequence parallelism alongside TP to keep activation memory down.
The file predates pre-commit being run on this branch and failed both hooks at HEAD. No functional change: import grouping per isort --profile black, plus black wrapping three call sites that exceeded the 110-column limit.
Two opt-in hooks that made the Gemma 4 verification and tuning work possible
without editing third_party/Megatron-Bridge or widening every recipe signature:
PRIMUS_GEMMA4_DUMP_MODEL=1 dump the instantiated module tree, the
per-layer attention geometry and parameter
totals, to check the built model against the
published architecture
PRIMUS_GEMMA4_PROFILE=<a>:<b> run the torch profiler over a step range
PRIMUS_GEMMA4_SET="a.b=c;..." assign arbitrary dotted ConfigContainer
fields after the recipe builds it
The recipes expose a fixed keyword list, and Primus' flat YAML overrides reach
only top-level ConfigContainer attributes, so the knobs that matter most for
Gemma 4 -- recompute_granularity, moe_token_dispatcher_type,
expert_tensor_parallel_size, cp_comm_type -- were otherwise unreachable. This is
a tuning hook rather than a supported surface; settings worth keeping belong in
the recipe.
Bridge implements Gemma 4's RMSNorm as the literal Hugging Face expression, which promotes the activation to fp32 twice and leaves the square / mean / rsqrt / scale chain unfused. Every norm in both the dense and MoE stacks routes through it, so a profiled dense step spends a large share of its GPU time in elementwise kernels that are almost entirely memory traffic. PRIMUS_GEMMA4_FUSED_NORMS=compile hands the same arithmetic to torch.compile so Inductor collapses the chain into a single kernel that reads and writes the activation once. =te routes the scaled norms to Transformer Engine's fused RMSNorm instead; that is not bit-comparable with Hugging Face, since TE accumulates in fp32 rather than promoting the whole expression. Dropping the fp32 temporaries also frees activation memory, which raises the reachable micro-batch size. Default stays eager.
Throughput on both variants is set almost entirely by how large a micro-batch fits, so record which opt-in patch or override frees the memory to raise it: the TE attention and fused-norm patches on the dense side, and activation recompute on the MoE side. Also note that TP=1/EP=8 is the faster of the two MoE layouts where it fits, since the expert all-to-all is latency-bound rather than bandwidth-bound. Defaults are unchanged; micro_batch_size stays at 1 so the shipped configs run without any of these enabled.
Gemma4ModelProvider accepts transformer_impl but never acts on it: the layer spec is selected from HAVE_TE before the field is read, so a user who asks for transformer_impl="local" silently gets TE layers and no warning. Upstream's own E4B dense recipe sets transformer_impl="local", which means that recipe cannot currently be reproduced through this provider. Fix it in a patch rather than in the third_party submodule, matching the approach already used for gemma4.dense_attention. Also adds PRIMUS_GEMMA4_TORCH_OPTIM=1, which points the optimizer at torch instead of TE's FusedAdam. This is deliberately independent of the layer spec: on gfx1250 TE's FusedAdam does not return -- optimizer.step() stays pinned at full GPU utilisation with an unchanging stack in multi_tensor_apply, while the driver logs "MES(0, 0) ring buffer is full". Without an escape hatch that works while TE layers are built, testing TE on that device risks a host reboot.
…type Several of the fields worth overriding are enums rather than strings, most importantly model.attention_backend, which is an AttnBackend. setattr-ing the plain string "unfused" onto it does not raise -- megatron-core compares the field by enum identity, so the match silently fails and the run proceeds on the default backend. That is the worst outcome for a tuning hook: a knob that reports success and changes nothing, so an A/B comparison returns two identical numbers and reads as "this setting makes no difference". Coerce by enum name first, then by value, and leave the parsed value untouched if neither resolves so genuinely bad input still surfaces.
… directory
Some ROCm packages install hipBLASLt's Tensile files one directory below where
the library loads them: the kernels land in
<prefix>/hipblaslt/library/gfx<arch>/TensileLibrary_lazy_gfx<arch>.dat
while hipBLASLt looks for
<prefix>/hipblaslt/library/TensileLibrary_lazy_gfx<arch>.dat
The resulting failure is quiet and expensive. The load fails on stderr, no
solution is ever found, and hipblasLtMatmulAlgoGetHeuristic then returns
HIPBLAS_STATUS_INVALID_VALUE for *every* shape. hipBLASLt therefore looks
completely broken rather than merely untuned, and the natural response -- set
TORCH_BLAS_PREFER_HIPBLASLT=0 and accept the rocBLAS fallback -- is a large
silent performance loss. On a gfx1250 wheel, setting this variable takes a plain
square bf16 GEMM from failing outright to working and speeds up single-GPU
Gemma 4 training substantially, with the loss curve unchanged. It also
un-blocked Transformer Engine, which had appeared broken for the same reason.
Detected rather than hardcoded, and deliberately conservative:
* It is a no-op unless the top-level file is absent *and* a per-architecture
copy exists, so correct packaging is never second-guessed.
* An existing HIPBLASLT_TENSILE_LIBPATH is left alone.
* If several architectures are shipped side by side it only acts on a match it
can confirm via offload-arch, rather than guessing.
Placed in base_env.sh rather than a per-accelerator file because it is sourced
unconditionally: GPU detection reports "unknown" on some of the affected
hardware, so an MI<model>.sh would not be loaded on the very machines that need
this.
Candidate ordering is load-bearing. A ROCm pip install can ship two copies of
hipBLASLt whose Tensile trees differ -- an architecture-specific runtime wheel
and a development tree with fewer code objects -- and neither obvious selector
is reliable: ROCM_HOME/ROCM_PATH can name the development tree, and ldd on
libtorch_hip.so resolves libhipblaslt.so.1 there too, while /proc/self/maps
after importing torch shows the runtime wheel's copy mapped instead. The
architecture-specific runtime wheel is therefore preferred, being both the copy
observed in use and the more complete one.
Neither model fits on one MI455X with a full optimizer state, so these are layer-reduced proxies: num_layers drops to 6 (preserving the published 5 sliding : 1 full attention ratio) while every other architectural dimension stays as published, so per-layer cost remains representative. Both configs record the settings that were measured fastest, the two environment variables that are not optional on gfx1250 (HIPBLASLT_TENSILE_LIBPATH, without which hipBLASLt fails every solution lookup; PRIMUS_GEMMA4_TORCH_OPTIM for TE), the fusions that have no gfx1250 build, and a warning against quoting the harness TFLOP/s figure for a layer-reduced model, since it derives FLOPs from the configured depth rather than the built one. The 31B config also records that its ceiling is stability rather than memory: raising micro_batch_size or num_layers past the values here has repeatedly wedged the driver well below the memory limit.
The branch ships no tests, which for an architecture with this many unusual features is the main gap in reviewing it. These are the cheapest ones worth having: the override hook is pure Python, so the whole file runs in well under a second with no megatron-core, no GPU and no checkpoint. Twelve cases over three helpers. The enum ones are the reason the file exists -- they guard the silent failure where assigning the string "unfused" to model.attention_backend does not raise, fails to match megatron-core's enum identity, and leaves the run on the default backend. A tuning hook that reports success and changes nothing turns an A/B comparison into two identical numbers, which reads as "this setting does not matter" rather than "this setting was never applied". Two of the tests assert on log output rather than return values, because for a dropped override being reported matters as much as not being applied: a typo that vanished quietly would leave the run on defaults while the operator believed otherwise. log_rank_0 is captured via monkeypatch. Primus's logger is process-global and only initialised by a real run, so calling it from a bare unit test raises AttributeError on a None logger; capturing rather than suppressing is what lets those two tests check the message. A local Enum stands in for AttnBackend to keep the file free of megatron-core, which matters more than it looks: importing the Gemma 4 recipe pulls in megatron.bridge.recipes, whose __init__ eagerly imports every recipe family including flux, so it needs the diffusion stack stubbed or installed to load.
…izer switch Tests for the silent failure the local_spec patch exists to fix: the provider binds its layer spec from a module-level HAVE_TE at dataclass-construction time and never consults transformer_impl, so asking for local layers leaves a config that contradicts itself and quietly builds TE layers anyway. Eleven cases driven against _rebind_spec with stub containers, so the file needs no megatron-core and runs in well under a second. Two of them cover properties a naive fix would get wrong: a config that never asked for local layers must be left exactly alone, and the PRIMUS_GEMMA4_TORCH_OPTIM switch must work independently of the layer spec, since the runs that need it are exactly the TE runs the rest of the function declines to touch. Writing them surfaced a real ordering bug, now fixed. _use_torch_optimizer returned early when megatron.core.optimizer was not yet imported, and that early return also skipped clearing use_precision_aware_optimizer -- a config field that needs only the container. So whether the field got cleared depended on patch-vs-import ordering that this patch does not control, and the bad outcome was a half-applied switch: torch Adam requested while precision-aware optimisation stayed on, relying on FusedAdam master weights that are no longer there. The two actions are now independent, and the import-ordering case says so in the log instead of returning silently.
… FusedAdam
CPU optimizer offload is the only way to keep optimizer state off the device, but
megatron's offload path always builds its GPU-side optimizer from TE's FusedAdam,
and neither documented way of avoiding that works:
* use_torch_optimizer_for_cpu_offload is inert for optimizer='adam'. The offload
branch honours it at optimizer/__init__.py:514 and then unconditionally
overwrites the result at line 516. Since use_precision_aware_optimizer already
asserts adam, the field never takes effect in practice.
* USING_PYTORCH_OPTIMIZER, which gemma4.local_spec sets, is read only at line 556
inside the non-offload adam branch. The offload path never consults it, so
enabling offload silently undoes that override.
Add an opt-in gemma4.cpu_offload patch (PRIMUS_GEMMA4_CPU_OFFLOAD=<fraction>) that
configures offload and rebinds the module-level Adam name that line 516 actually
reads. That name has two readers -- line 516, and line 560 which is unreachable once
USING_PYTORCH_OPTIMIZER is true -- and no isinstance checks, so the rebind is
contained. The CPU half was already torch.optim.AdamW, so both halves of the hybrid
optimizer now run the same update rule.
Precision-aware optimisation is kept on here, which is the opposite of what
gemma4.local_spec does, because the two cases genuinely differ: local_spec clears it
since a bare torch Adam has no master weights, whereas HybridDeviceOptimizer is
constructed with param_update_in_fp32=True and maintains its own fp32 copies plus the
decoupled_grad plumbing. megatron agrees -- optimizer_config.py:444 returns early from
the "precision-aware requires TE FusedAdam" check when offload is set. Clearing it
would put the fp32 master weights back on the device, which is the memory offload
exists to reclaim.
Because the two patches disagree about that one field, local_spec now asks
gemma4_cpu_offload.offload_fraction() rather than inspecting the config, so the
outcome does not depend on which patch is applied first.
Tested: 15 tests covering fraction parsing, the Adam rebind and its idempotence, the
fields the offload path asserts on, and the cross-patch coupling in both directions.
The local_spec patch now consults PRIMUS_GEMMA4_CPU_OFFLOAD to decide whether to leave use_precision_aware_optimizer alone, so an inherited value from the caller's environment would quietly invert what these tests assert. Clear it in the fixture alongside PRIMUS_GEMMA4_TORCH_OPTIM.
… patches gemma4_moe_post_attn_norm fixes a silent correctness bug. _gemma4_block_spec attaches the post-attention normalization only on the TransformerEngine branch, so the non-TE MoE path builds a model that is missing a normalization layer entirely. Nothing fails: the run starts, the loss descends, and it converges to the wrong thing. That makes it the most dangerous defect in this series, because every other issue here announces itself. It is also architecture-independent -- nothing about it is specific to gfx1250 -- and the fix is small, since the norm already exists in the spec and only needs attaching outside the TE guard. Corroborating evidence that this is MoE-only: on the 31B dense model, loss agrees within 0.040 nats across two layer impls and two GEMM backends, and the absence of any gap there is what localises the defect to the MoE spec. gemma4_attn_backend selects the attention backend, which matters on gfx1250 because TE's fused attention is not built in this image and the fallback has to be chosen deliberately rather than stumbled into. gemma4_chunked_logits chunks the vocabulary projection. On a layer-reduced proxy the projection is a large share of both memory and FLOPs -- 32% on the dense model and 58% on the MoE -- so chunking it is worth having available. Both are opt-in and no-ops when unset. All three register through @register_patch like the existing Gemma 4 patches.
…does not apply A skipped override used to be a log line at INFO and nothing more, which is the worst outcome available for a tuning hook: a mistyped field name is discarded, the run succeeds, and it reports a clean number for a configuration nobody asked for. That result is indistinguishable from a real one unless someone happens to read the right log line, and an experiment is worthless if the knob under test may silently not have been set. This collects the overrides that did not land and refuses to start instead. PRIMUS_GEMMA4_SET_LENIENT=1 downgrades it back to a warning, which is useful when sharing one override string across configs whose field sets differ. This is the same failure shape as the enum coercion in the preceding commit -- both are cases where the config system reported success and changed nothing.
…eference numbers, add a gfx1250 runbook
Three things, all on the MI455X single-GPU Gemma 4 recipes.
Depth moves out of the example overrides into two new model presets,
gemma4_{26b,31b}_6layer_proxy.yaml, so the example configs carry no architecture
parameters. Note that a Megatron-Bridge preset is a recipe pointer
(recipe/flavor/hf_path) rather than an architecture spec, so depth has to be a
nested `model:` block: a flat top-level num_layers is neither a recipe kwarg nor
a top-level ConfigContainer field and is dropped in silence. parser.py keeps a
preset's own `model` block instead of replacing it with the preset name, which is
what makes this work. Verified two ways -- the fully merged config is identical
to the previous inline form, and on hardware the run logs
"Set config_container.model.num_layers = 6".
The reference numbers in both configs were wrong, and wrong in a way worth
spelling out. The 31B header asserted that micro_batch_size 8, num_layers 12,
and both TransformerEngine attempts had "wedged the driver", and concluded a
stability ceiling of micro_batch_size 4. Every one of those four claims rested on
a single wedge, and all four have since been overturned by re-probing on a
healthy boot: mbs 8 completes twice at 14636 tok/s, L12 and L18 are clean, and TE
runs fine. On this device a wedge is evidence about the boot, not about the
configuration. Only micro_batch_size 16 has failed twice and is the real ceiling.
Both reference lines were also internally inconsistent, having been assembled
from runs at different micro-batches: the 31B quoted 760 ms/iter with 4300 tok/s
and 149.3 GB, which are three different configurations (760 ms is mbs 1, 149.3 GB
is mbs 4, and 760 ms at mbs 4 would be 10779 tok/s). The 26B paired 558 ms/iter,
which is its mbs 1 figure, with the mbs 16 throughput and memory. Both now quote
precisely-labelled rows with two clean runs each, and state plainly that the
shipped defaults do not reach the FlyDSL numbers because that routing is not part
of this repo.
Finally, a runbook at docs/04-technical-guides/Gemma4_Bridge_MI455X_README.md,
following the structure of the native LoRA MI455X guide and wired into the same
five index and toc files. It documents the hardware gate, the hipBLASLt Tensile
path fix and why TORCH_BLAS_PREFER_HIPBLASLT=0 is the wrong reaction to it, the
TE FusedAdam guard, what a passing run prints, and the intermittent MES wedge
including the ten-minute window in which a wedged GPU still reads as healthy.
yeandy
force-pushed
the
dev/implement-gemma-megatron
branch
from
September 15, 2026 19:05
1fc221d to
7b34418
Compare
There was a problem hiding this comment.
Code Review: AMD-AGI/Primus
Summary
1 architecture concern, 34 files reviewed.
Rules checked: 7 (7 passed, 0 failed).
Architecture Concerns
These findings come from architecture principles distilled from the repo's PR lineage — patterns that caused PRs to be rejected or significantly reworked.
batch-related-dependency-bumps-atomically
When multiple dependencies belong to the same ecosystem/manifest and are managed by an automated dependency-update tool that groups them (e.g., a dependency group across a single directory), they should be bumped together in one coordinated PR rather than as a scattered series of one-off single-package PRs. Piecemeal single-package bumps each trigger a full CI/review cycle, leave the manifest in inconsistent intermediate states where interdependent packages are mismatched, and multiply merge-conflict and re-validation surface across the same set of files.
Findings:
- This PR advances the
third_party/Megatron-Bridgesubmodule (9577b12 -> 51885cf, i.e. v0.6.1) but leaves its coupled Python pin untouched:runner/helpers/hooks/train/pretrain/megatron_bridge/requirements-megatron_bridge.txtstill saystransformers==5.10.1while the new comment states Bridge v0.6.1 requires>=5.8,<=5.12.1and that Gemma 4 was actually measured on 5.12.1. Bump the submodule and thetransformerspin (to the version you validated) in the same commit so the Megatron-Bridge dependency group moves atomically.
Posted by PR Pundit — AI-powered code review grounded in this repo's merge history.
…withhold absolute throughput Two things: a measurement the configs were missing, and a decision about what to publish for this part. The 31B config had no reference for the configuration it actually ships. Every figure we held was at a different micro-batch, a different recompute setting, or used FlyDSL routing that is not in this repo, so its header could only bracket. Measured it directly -- micro-batch 4, recompute none, local layers, no FlyDSL, seq 2048 -- with two clean runs at 0 nan and 0 skipped iterations. Run-to-run agreement is 0.48%, peak memory is identical to 0.01 GB, and the loss is identical to eight significant figures. **Absolute throughput for MI455X is deliberately not recorded.** This repository publishes tok/s and ms/iter freely for MI300X and MI355X but nowhere for MI455X: the native LoRA MI455X guide states outright that "proxy numbers are enablement signals, not model TFLOP/s" and ships four recipes without a single throughput figure, and the only MI455X performance constants anywhere are the projection tool's *estimated* hardware peaks, which its own design notes describe as not officially published. Rather than make these the first measured MI455X throughput numbers in the tree, the configs and runbook now give peak VRAM and loss as the fingerprint to check a run against, plus the speedup ratios. That loses very little and gains accuracy. The ratios are what make a reader act -- 2.93x from the hipBLASLt Tensile path, 4.73x from FlyDSL at this shape against 3.13x at micro-batch 1, 1.69x for TE on the MoE, an 8% regression for FlyDSL on the MoE -- and unlike a tok/s figure they survive a change of image or host, which matters here because we have a case of the same configuration measuring 27% apart from an image version alone. For "is my run healthy", peak memory and loss are the stricter test anyway, since both reproduce exactly on this setup while throughput does not. Peak memory matching the recorded FlyDSL row at the same micro-batch to 0.03 GB confirms those two differ only in GEMM routing, which is what makes the 4.73x exact rather than approximate. The loss agreeing to four significant figures across the two GEMM backends is correctness evidence for that routing at a second shape. Micro-batch 8 is now explicitly not recommended on the unpatched path. It is clean twice with FlyDSL, but without it has wedged twice out of two attempts, and one of those was minutes after two clean micro-batch 4 runs on the same boot. That same-boot pairing is the only wedge evidence in this work where boot state is held fixed -- every previous "configuration X wedges" claim collapsed precisely because the clean and wedged observations came from different boots -- so it is recorded as an interaction between micro-batch and the GEMM path rather than as a ceiling. The runbook also gains a "See also" section for the other two MI455X efforts in the repo, because a reader on this hardware will want them and they are easy to miss. It records what transfers and what does not: the native LoRA MI455X guide needs PRIMUS_TURBO_ATTN_BACKEND=triton where the Bridge path does not. The DeepSeek-V4 gfx1250 launcher is the source of several environment settings that base_env.sh and MI455X.sh now apply automatically, and it documents the same unrecoverable-MES wedge from four unrelated triggers -- including one that presented with a completely clean dmesg, which is why this guide says no driver-log signal is sufficient to detect it. Two gfx1250 facts from that work generalise and are now recorded: Primus-Turbo's gluon and FlyDSL attention kernels are gfx950-only, and all_reduce with op=AVG hangs on some gfx1250 builds even at world size 1, which matters for MoE auxiliary-loss reductions.
yeandy
force-pushed
the
dev/implement-gemma-megatron
branch
from
September 15, 2026 19:41
7b34418 to
c40dca7
Compare
There was a problem hiding this comment.
Code Review: AMD-AGI/Primus
Summary
No violations found, 34 files reviewed.
Rules checked: 7 (7 passed, 0 failed).
✅ Passed
All 7 rules satisfied. No architecture concerns found.
Posted by PR Pundit — AI-powered code review grounded in this repo's merge history.
…ich transformers version Went to verify the two reference fingerprints against the tree as pushed and found they do not have equal standing, so the table now says so instead of presenting them as equivalent. The 31B row was measured on this revision, after the rebase. The 26B row was measured four days and 119 upstream commits earlier and has not been re-measured since. The configuration itself did not change, but the code underneath it did, so it is now labelled as indicative rather than as a guarantee, with a request to report a divergent result rather than assume the table is right. Also recorded which transformers version the numbers are actually from, which turns out not to be the one the repo ships. The pretrain hook's requirements-megatron_bridge.txt pins 5.10.1, but the container image we measure in carries 5.12.1, and `primus-cli direct` bypasses the hook that would install the pin -- so every number taken so far, before and after the rebase, is on 5.12.1 and the shipped pin has never been executed. Both versions are inside Megatron-Bridge v0.6.1's >=5.8,<=5.12.1 range, so this is a provenance caveat rather than a defect, but a reader going through the normal runner path is on a version these numbers were not taken on and should know that.
There was a problem hiding this comment.
Code Review: AMD-AGI/Primus
Summary
1 architecture concern, 34 files reviewed.
Rules checked: 7 (7 passed, 0 failed).
Architecture Concerns
These findings come from architecture principles distilled from the repo's PR lineage — patterns that caused PRs to be rejected or significantly reworked.
batch-related-dependency-bumps-atomically
When multiple dependencies belong to the same ecosystem/manifest and are managed by an automated dependency-update tool that groups them (e.g., a dependency group across a single directory), they should be bumped together in one coordinated PR rather than as a scattered series of one-off single-package PRs. Piecemeal single-package bumps each trigger a full CI/review cycle, leave the manifest in inconsistent intermediate states where interdependent packages are mismatched, and multiply merge-conflict and re-validation surface across the same set of files.
Findings:
- This PR advances the
third_party/Megatron-Bridgesubmodule (9577b128 -> 51885cf1, described in the docs as v0.6.1) but only adds a comment torunner/helpers/hooks/train/pretrain/megatron_bridge/requirements-megatron_bridge.txtwithout touching any pin. That requirements file is explicitly 'Extracted from third_party/Megatron-Bridge/pyproject.toml', so it is the same manifest/ecosystem as the submodule: re-sync every extracted pin (transformers and the 'Core dependencies from [project.dependencies]' block) against the new submodule commit in this same PR rather than letting them drift. - The
transformers==5.10.1pin is now inconsistent with the rest of the change: the new docs and both MI455X configs record all reference numbers as measured on transformers 5.12.1, and the added comment concedes 'primus-cli direct bypasses the hook that installs the pin, so if you go through the normal runner path you are on 5.10.1 and these numbers were not taken there'. Bumptransformersto the version actually validated (5.12.1, still inside Megatron-Bridge v0.6.1's >=5.8,<=5.12.1 range) in the same commit as the submodule bump, so the runner path and the documented fingerprints do not diverge.
Posted by PR Pundit — AI-powered code review grounded in this repo's merge history.
…e shipped transformers pin Re-measured the 26B config on the current tree and tested the version the repo actually pins, which resolves the two caveats added in the previous commit. The 26B fingerprint is now measured on this revision. Peak memory came back at 301.58 GB against the 301.6 GB recorded before the rebase -- unchanged to the hundredth of a gigabyte -- with 0 nan and 0 skipped iterations. The loss moved slightly, 15.049 to 15.048, with the configuration untouched, so the table now quotes the measured 15.048. That small movement is the useful part, and the surrounding text has been rewritten around it. Peak memory has now held to 0.01 GB across repeats, across a rebase onto a much newer main, and across a change of transformers version, while loss is stable within a revision but can move in the third decimal when the code underneath changes. So the guidance is no longer "both reproduce exactly": memory is the reliable half, a third-decimal loss difference means the code moved rather than that the reader's run is broken, and a first-decimal difference is a real problem. The transformers pin is now tested rather than just flagged. On an image identical to ours except for transformers downgraded to the pinned 5.10.1, the 31B reproduced both of its numbers exactly -- 149.3 GB and loss 26.056 -- so that fingerprint is version-independent across Megatron-Bridge v0.6.1's supported range. The 26B attempt on 5.10.1 lost the device to the MES wedge after the model had built and entered the training loop, so it stays marked as unconfirmed, with the runbook saying plainly that nothing implicates 5.10.1 there. Worth stating because the same-boot sequence looks incriminating out of context: the 26B ran clean on 5.12.1 and then wedged on 5.10.1 two minutes later. But that same configuration has also wedged *on* 5.12.1 and run clean on it, so the difference is the boot lottery, not the version.
There was a problem hiding this comment.
Code Review: AMD-AGI/Primus
Summary
1 architecture concern, 34 files reviewed.
Rules checked: 7 (7 passed, 0 failed).
Architecture Concerns
These findings come from architecture principles distilled from the repo's PR lineage — patterns that caused PRs to be rejected or significantly reworked.
batch-related-dependency-bumps-atomically
When multiple dependencies belong to the same ecosystem/manifest and are managed by an automated dependency-update tool that groups them (e.g., a dependency group across a single directory), they should be bumped together in one coordinated PR rather than as a scattered series of one-off single-package PRs. Piecemeal single-package bumps each trigger a full CI/review cycle, leave the manifest in inconsistent intermediate states where interdependent packages are mismatched, and multiply merge-conflict and re-validation surface across the same set of files.
Findings:
- This PR advances third_party/Megatron-Bridge (9577b128 -> 51885cf1, documented as v0.6.1) but leaves the sibling pin in runner/helpers/hooks/train/pretrain/megatron_bridge/requirements-megatron_bridge.txt at transformers==5.10.1 with only a comment change. Bump the submodule and the transformers pin together in one atomic Megatron-Bridge dependency-group change so the vendored bridge and its declared HF dependency never sit at mismatched revisions.
- The docs added here state the reference numbers were measured on transformers 5.12.1 (the container image's version) while the hook installs 5.10.1, and that only the 31B fingerprint was confirmed on 5.10.1. Resolve this in the same PR: either raise the pin to the validated 5.12.1 alongside the Megatron-Bridge bump, or re-validate both configs on 5.10.1 — do not merge a state where the pinned version and the validated version diverge.
- Treat requirements-megatron_bridge.txt as a grouped manifest: if the Megatron-Bridge v0.6.1 bump changes any other constraint extracted from its pyproject.toml (beyond transformers), update all of those entries in this same PR rather than trickling them through separate follow-ups, each of which costs another CI + review cycle and leaves the hook's requirements partially stale.
Posted by PR Pundit — AI-powered code review grounded in this repo's merge history.
…reproduce The runbook and the 31B config published 4.73x / 3.13x for routing the dense linears through a Primus-Turbo FlyDSL kernel, and 0.92x for the same change on the 26B. Both texts noted the routing was "not yet in this repo", but they printed a bold multiple next to an environment variable name, which invites a reader to set it and conclude our numbers are wrong. They would be right to. PRIMUS_GEMMA4_TURBO_LINEAR is a silent no-op on this tree: the patch module is not listed in patches/gemma4/__init__.py, which is what imports the patch modules, so it never loads and setting the variable fails quietly instead of erroring. Two 31B mbs 4 runs on one boot, identical but for the variable, came out 2846.3 ms and 2790.8 ms -- a 2% spread, this host's noise, and both matching the unpatched baseline. The 4.73x was also cross-revision, dividing a measurement taken before a 119-commit rebase by one taken after, on a branch where the 26B loss moved across that same rebase. So the ratios come out of both files. The mechanism stays, since it is the useful part and is independently supported by the hipBLASLt Tensile row: exactly one of four bf16 contraction layouts is tuned on gfx1250, so dgrad and wgrad run roughly 15x slower than the tuned one. It is now stated without a number, with the no-op called out explicitly, and the figure deferred to the PR that actually ships the routing. Every ratio the runbook still prints is reproducible from this tree, which it now says. Also retires the last cross-revision claim on the branch, and corrects the transformers note: the 26B is unconfirmed on the pinned 5.10.1 after two attempts on two boots, not merely untested. Both lost the device to the MES wedge after entering the training loop. That row is 0 of 2 on 5.10.1 against 1 of 1 on 5.12.1, which we are not reading as version causation -- the same-boot control in the ticket already showed this configuration wedging on 5.12.1, and at this host's failure rate two consecutive losses on one row is unremarkable -- so the fingerprint is simply described as established on 5.12.1 only.
There was a problem hiding this comment.
Code Review: AMD-AGI/Primus
Summary
1 architecture concern, 34 files reviewed.
Rules checked: 7 (7 passed, 0 failed).
Architecture Concerns
These findings come from architecture principles distilled from the repo's PR lineage — patterns that caused PRs to be rejected or significantly reworked.
batch-related-dependency-bumps-atomically
When multiple dependencies belong to the same ecosystem/manifest and are managed by an automated dependency-update tool that groups them (e.g., a dependency group across a single directory), they should be bumped together in one coordinated PR rather than as a scattered series of one-off single-package PRs. Piecemeal single-package bumps each trigger a full CI/review cycle, leave the manifest in inconsistent intermediate states where interdependent packages are mismatched, and multiply merge-conflict and re-validation surface across the same set of files.
Findings:
- This PR bumps the
third_party/Megatron-Bridgesubmodule (9577b12 -> 51885cf) but leaves every pin inrunner/helpers/hooks/train/pretrain/megatron_bridge/requirements-megatron_bridge.txtuntouched — that file's own header says it is 'Extracted from third_party/Megatron-Bridge/pyproject.toml'. Re-derive and bump the whole extracted requirement set in the same PR so the submodule and its mirrored pins move atomically instead of drifting. transformers==5.10.1is left stale while the diff only adds a comment admitting the validated/container version is 5.12.1 and that Megatron-Bridge v0.6.1 allows >=5.8,<=5.12.1. Bump the pin to the version actually validated (5.12.1) in this same PR rather than documenting the mismatch; shipping a comment that contradicts the pin guarantees a second review/CI cycle for the follow-up bump.- The Megatron-Bridge submodule bump and the megatron_bridge hook requirements belong to one coordinated update group. Consolidate them into a single grouped bump PR (submodule + transformers + any other pins whose upstream ranges moved with v0.6.1) so there is no intermediate state where the backend code targets a newer Bridge revision than its requirements file describes.
Posted by PR Pundit — AI-powered code review grounded in this repo's merge history.
Re-measuring the 2.93x hipBLASLt Tensile ratio on the shipped tree did not produce a corrected ratio. It produced the finding that there is no ratio to correct: without HIPBLASLT_TENSILE_LIBPATH pointed at the directory that holds the files, these configs abort on the first GEMM with HIPBLAS_STATUS_INVALID_VALUE out of hipblasLtMatmulAlgoGetHeuristic. Three runs on one boot, 31B at mbs 4, device healthy throughout: the corrected path completed 8 iterations and agreed with the four step times taken earlier the same day to within 2%, and the broken path failed twice out of two attempts with zero iterations. The broken arm was reached by pre-setting the variable, which the detector honours because it is guarded on the variable being unset -- so this also exercises that guard. PROGRESS 4.36 saw 2803 ms on the broken arm because the failing heuristic call was absorbed by a fallback. On this revision, at this config, with torch's default TORCH_BLAS_PREFER_HIPBLASLT (the repo never sets it), the error propagates and kills the run. That makes the detector the difference between training and not training, which is a stronger justification than a speedup but a different claim, so the runbook and both configs now state it as one. The Tensile row comes out of the step-time table, because a works/does-not-work gate does not belong in a column headed "effect on step time". The 2.93x survives only as an explicit note that it has been superseded, so nobody later mistakes it for a current figure. base_env.sh needed no edit: it already said "from failing outright to working" and quoted no ratio. The four remaining ratios are now dated. They are sound as ratios -- each is a same-session A/B with both arms measured back to back, unlike the FlyDSL figure removed in the previous commit -- but they predate this branch's rebase and were taken at layer and window configurations that are not the two shipped here, so the runbook now vouches for their direction and rough magnitude and explicitly not for the second digit.
There was a problem hiding this comment.
Code Review: AMD-AGI/Primus
Summary
1 architecture concern, 34 files reviewed.
Rules checked: 7 (7 passed, 0 failed).
Architecture Concerns
These findings come from architecture principles distilled from the repo's PR lineage — patterns that caused PRs to be rejected or significantly reworked.
batch-related-dependency-bumps-atomically
When multiple dependencies belong to the same ecosystem/manifest and are managed by an automated dependency-update tool that groups them (e.g., a dependency group across a single directory), they should be bumped together in one coordinated PR rather than as a scattered series of one-off single-package PRs. Piecemeal single-package bumps each trigger a full CI/review cycle, leave the manifest in inconsistent intermediate states where interdependent packages are mismatched, and multiply merge-conflict and re-validation surface across the same set of files.
Findings:
- This PR advances the third_party/Megatron-Bridge submodule (to the v0.6.1-era commit that adds Gemma 4 support) but leaves runner/helpers/hooks/train/pretrain/megatron_bridge/requirements-megatron_bridge.txt pinned at transformers==5.10.1 while only adding explanatory comments. Bump the transformers pin in the same PR as the submodule move (to the version the Gemma 4 work was actually validated on, 5.12.1, within Bridge's >=5.8,<=5.12.1 range) so the Megatron-Bridge dependency group lands atomically instead of leaving the manifest stale relative to the submodule.
- Avoid shipping a documented pin/runtime mismatch as an intermediate state: the new Gemma4_Bridge_MI455X_README.md and the MI455X configs record fingerprints measured on transformers 5.12.1 (container image) while the hook installs 5.10.1, and explicitly note the 26B row is unconfirmed on 5.10.1. Either update the pin together with the submodule bump in this PR, or split the dependency changes out into one coordinated Megatron-Bridge group-bump PR (submodule + transformers + any other Bridge-derived entries in requirements-megatron_bridge.txt) so the pin and the submodule are reviewed/CI'd once, not across two cycles.
Posted by PR Pundit — AI-powered code review grounded in this repo's merge history.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.