Skip to content

[codex] Port Gemma4 PR1855 training fixes#2

Draft
EazyReal wants to merge 1 commit into
leofan-lab:gemma4-prfrom
EazyReal:codex/gemma4-pr1855-training-fixes-20260618
Draft

[codex] Port Gemma4 PR1855 training fixes#2
EazyReal wants to merge 1 commit into
leofan-lab:gemma4-prfrom
EazyReal:codex/gemma4-pr1855-training-fixes-20260618

Conversation

@EazyReal

Copy link
Copy Markdown

What changed

Ports the upstream-shaped Gemma4 training fixes from the internal v031 integration work onto PR1855 without pulling in Vmax agent/launcher/retry glue.

  • Gemma4 bridge QKV import now chooses packed-QKV geometry from the HF tensors being imported: [q, k] means global K=V geometry, [q, k, v] means sliding/local geometry. This covers pipeline-parallel imports where a local Megatron layer name can otherwise look like layer 0.
  • Gemma4 provider softcap now uses an in-place custom autograd function so final logit softcapping does not allocate a second logits tensor while preserving the tanh derivative.
  • Adds focused CPU tests for the PP QKV import invariant, softcap storage/gradient behavior, and packed RoPE kwargs forwarding.

Why

The internal bringup found that PR1855’s model path was mostly structurally complete, but training/export still needed narrow fixes around QKV import shape selection and softcap memory pressure. The other internal commits named in the audit map to behavior already present in PR1855’s larger native implementation: concatenated dual-RoPE with packed kwargs, explicit pre-MLP norm, dense MLP without padding-mask forwarding, SDPA/flash CP attention, and config-local layer shape overrides.

Validation

Passed:

  • uv run --python 3.12 --with ruff ruff check slime_plugins/mbridge/gemma4.py slime_plugins/models/gemma4_provider.py tests/gemma4/test_gemma4_qkv_roundtrip.py tests/gemma4/test_gemma4_provider.py tests/gemma4/test_gemma4_dual_rope.py
  • git diff --check
  • uv run --python 3.12 --with torch --with numpy python -m py_compile slime_plugins/mbridge/gemma4.py slime_plugins/models/gemma4_provider.py tests/gemma4/test_gemma4_qkv_roundtrip.py tests/gemma4/test_gemma4_provider.py tests/gemma4/test_gemma4_dual_rope.py
  • Stubbed runtime checks for the changed logic: QKV PP/global shape selection, in-place softcap gradient, and DualRotaryEmbedding packed kwargs forwarding.

Could not run full targeted pytest locally because this checkout environment does not have external mbridge or megatron packages available; collection stops before executing tests with ModuleNotFoundError: No module named 'mbridge' / No module named 'megatron'.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant