Gemma 4 31B random-init pretraining shows transient loss spikes in both AutoModel and Megatron Bridge

Description

We observed a large transient loss spike while comparing Gemma 4 31B random-initialized pretraining between NeMo AutoModel and Megatron Bridge.

The spike is present in both implementations under the same data, optimizer, LR schedule, batch size, sequence length, seed, and parallel topology. It also persists after changing the AutoModel attention backend, changing the Megatron output-projection initialization, lowering the LR to a constant 5e-5, and keeping numerically sensitive Megatron operations in FP32.

This issue is intended to track a shared Gemma 4 training-instability problem rather than attribute it to either framework.

Primary W&B evidence

Full comparison report:

• AutoModel vs Megatron Bridge convergence overlay

Implementation	W&B run	Peak optimizer step	Peak loss	Final loss
AutoModel, SDPA	19mnq6cb	18	17.0875	6.4423
AutoModel, FFPA/Flex	7jab6p0r	17	16.7645	6.3350
Megatron Bridge, default residual-scaled output init	ftsox54u	23	21.9531	6.6248
Megatron Bridge, all output projections initialized with std 0.02	4i3mvzts	66	19.7420	7.5609

All four runs completed 100 optimizer steps with finite losses and no skipped or NaN iterations.

AutoModel logs native steps starting from zero; the optimizer-step numbers above are normalized to start from one.

Spike details

AutoModel SDPA:

optimizer step 16: loss=13.9786, grad_norm=1133.81, lr=1.20e-4
optimizer step 17: loss=13.5424, grad_norm=37.12,   lr=1.275e-4
optimizer step 18: loss=17.0875, grad_norm=35.79,   lr=1.35e-4  <-- spike
optimizer step 19: loss=14.2678, grad_norm=29.63,   lr=1.425e-4

AutoModel FFPA/Flex:

optimizer step 15: loss=13.2473, grad_norm=295133.13, lr=1.125e-4
optimizer step 16: loss=11.6565, grad_norm=6597.32,   lr=1.20e-4
optimizer step 17: loss=16.7645, grad_norm=418.02,    lr=1.275e-4 <-- spike
optimizer step 18: loss=14.5337, grad_norm=135.74,    lr=1.35e-4

Megatron Bridge:

optimizer step 20: loss=12.7288, grad_norm=171291.41, lr=1.50e-4
optimizer step 21: loss=14.2880, grad_norm=11438.22,  lr=1.575e-4
optimizer step 22: loss=14.1685, grad_norm=1123.99,   lr=1.65e-4
optimizer step 23: loss=21.9531, grad_norm=50.60,     lr=1.725e-4 <-- spike
optimizer step 24: loss=17.3045, grad_norm=154.22,    lr=1.80e-4

The spike is preceded by a several-orders-of-magnitude collapse in the logged pre-clipping gradient norm in both frameworks.

The exact peak step differs because the two implementations are not expected to be numerically identical, but the qualitative sequence is similar:

very large initial gradient norm
    -> rapid gradient-norm collapse
    -> sudden loss increase
    -> recovery and continued convergence

Reproduction configuration

model: google/gemma-4-31B-it architecture
weights: random initialization
functional_parameters: 30,697,345,280

hardware:
  nodes: 16
  gpus_per_node: 8
  gpu: NVIDIA H100 80GB HBM3
  total_gpus: 128

parallelism:
  tensor_parallel: 4
  pipeline_parallel: 1
  context_parallel: 1
  expert_parallel: 1
  sequence_parallel: false

training:
  sequence_length: 4096
  micro_batch_size: 1
  global_batch_size: 1024
  steps: 100
  seed: 1234
  precision: BF16 compute
  activation_checkpointing: full

optimizer:
  type: AdamW
  beta1: 0.9
  beta2: 0.95
  eps: 1e-8
  weight_decay: 0.1
  grad_clip_norm: 1.0
  master_parameters: FP32
  gradients_and_reduction: FP32
  adam_moments: FP32

lr_schedule:
  warmup_steps: 40
  initial_lr: 0
  max_lr: 3e-4
  decay_style: cosine
  min_lr_at_step_100: 3e-5

gemma:
  vocabulary_size: 262144
  final_logit_softcapping: 30.0

Source revisions:

NeMo AutoModel: 70e738bbe99e3aa42e5ebc006834a37ab2d49890
Megatron Bridge: 573e088c9c6740082c39744e03dc5b009e730ed4
Megatron Core:   6513e3e23d6b5eda6a1c934990b15e804237732b

Both frameworks consume the same pre-tokenized RP2 indexed data. The dataset hashes and complete frozen runtime contracts are recorded in the native W&B runs.

Additional experiments

Attention implementation

Changing AutoModel attention from SDPA to audited FFPA/Flex did not eliminate the spike:

• SDPA: peak loss 17.0875 at optimizer step 18
• FFPA/Flex: peak loss 16.7645 at optimizer step 17

The FFPA run records zero SDPA/eager fallback. This suggests the spike is not specific to the attention kernel.

Output-projection initialization

Megatron normally initialized residual output projections using:

0.02 / sqrt(2 * 60) = 0.00182574

We reran Megatron with attention-output and MLP-down projections initialized using the HF-style unscaled std 0.02.

W&B: 4i3mvzts

The initialization audit confirms that the relevant projections were materialized with std approximately 0.02.

This did not eliminate the spike. It delayed it from optimizer step 23 to step 66:

step 64: loss=12.0215, grad_norm=1808.34
step 65: loss=16.3754, grad_norm=1397.46
step 66: loss=19.7420, grad_norm=260.17  <-- spike
step 67: loss=16.6501, grad_norm=531.96

Therefore, initialization scale affects the timing and magnitude but does not appear to be the sole root cause.

Constant LR

A Megatron run using constant LR 5e-5 also reproduced the instability before it was stopped at step 45:

• W&B: m3qwh8bo
• Peak loss: 18.9961
• Peak step: 28
• LR: 5e-5

This indicates that the spike is not exclusively caused by reaching the original warmup LR of approximately 1.3e-4–1.7e-4.

FP32-sensitive operations

We also reproduced the spike in a smaller Megatron run after keeping the following operations/state in FP32:

• optimizer master parameters and Adam moments
• gradients and gradient reduction
• RMSNorm and value norm accumulation
• residual stream
• attention softmax
• cross entropy and loss reduction
• BF16 reduced-precision matmul disabled

W&B: um1u4321

The run still reached loss 18.9027 at step 30. This makes a single BF16-sensitive operation less likely to be the complete explanation.

Expected behavior

With 40 warmup steps, gradient clipping, FP32 optimizer state, and no NaN or skipped iterations, we expected the loss to decrease without a large transient regression of this magnitude.

At minimum, we would like to identify whether this is:

1. expected behavior for random-initialized Gemma 4 at this scale;
2. an initialization or optimizer-recipe issue;
3. related to Gemma-specific embedding scaling, post-normalization, Q/K/value normalization, or final-logit softcapping;
4. caused by gradient clipping/Adam dynamics following the initial gradient-norm collapse; or
5. evidence of a shared semantic mismatch in the two implementations.

The runs eventually recover and converge, so this is not a NaN/crash issue. The concern is the reproducible transient instability across frameworks and attention implementations.