fix(fsdp): restore full-layer activation checkpointing for KV-shared models

[athitten]

What does this PR do ?

Fixes a 16% memory regression on the nightly Gemma4-4B benchmark by putting KV-shared
models (Gemma4 E2B/E4B) back on full-layer activation checkpointing.

The bug

Activation checkpointing can be applied two ways:

• whole block — wrap the entire decoder layer, so attention is recomputed in backward
• sub-module — wrap mlp and the norms individually

The sub-module path deliberately skips self_attn for KV-shared models.

PR #3513 added and not _has_kv_sharing to the whole-block branch. Gemma4 E4B has
num_kv_shared_layers: 18, so it started taking the sub-module path instead. Attention
activations across all 42 language layers stopped being recomputed and stayed resident:

Applied submodule activation checkpointing to 42 layers:
  {'mlp': 42, 'attention': 0, 'pre_norm': 42, 'post_norm': 42, 'mot': 0}
                ^^^^^^^^^^^^^

Nothing crashed and the loss did not move, so only the memory metric caught it.

Why the guard is not needed

The guard exists because recomputing attention could write K/V into the cache a second
time. That does not happen here: Gemma4 passes its shared layers _Gemma4KVShareHolder,
whose update() just returns what it was given. Replaying a block writes nothing.

Evidence

2x GB200, examples/vlm_finetune/gemma4/gemma4_4b_mock.yaml, activation checkpointing on:

	peak memory	loss / grad_norm
this PR's parent	87.37 GiB	6.0085 / 125.5791 …
before #3513	70.06 GiB	identical
this PR	70.06 GiB	identical

70.06 GiB matches the nightly's pre-regression value of 7.005e+01 exactly, and losses
and gradient norms agree to four decimals across all 8 steps. This was purely a memory
regression, not a numerics change.

Changelog

• parallelizer.py: drop the and not _has_kv_sharing guard so KV-shared models get
  whole-block checkpointing again, keeping fix(checkpoint): preserve FSDP2 mixed precision during recompute #3513's mixed-precision-safe wrapper.
• test_parallelizer.py: test_hf_native_candidate_with_kv_sharing_uses_full_layer_checkpointing
  now asserts whole-block wrapping. fix(checkpoint): preserve FSDP2 mixed precision during recompute #3513 had rewritten this test to assert the regressed
  behaviour, so it went green on the way in.
• New functional test L2_Parallelism_VLM_Gemma4_KVShared_AC (2 GPUs, builds a small
  E4B-shaped model in-process, nothing staged). It asserts every decoder block is
  checkpoint-wrapped, and that gradients match an un-checkpointed run. It lands in the
  parallelism folder added by test(pp): add gemma4 parallelism parity tests [1/4] #3529, so no CI wiring changes are needed.

Verified the new test fails on the regressed code:

AssertionError: KV-shared Gemma4 decoder blocks [0, 1, 2, 3] are not checkpoint-wrapped;
activation checkpointing fell back to sub-module wrapping, which excludes self_attn

and passes with the fix (OK (63 gradients compared)).

Peak-memory magnitude stays with the nightly benchmark — a proxy this small can't measure
it without a threshold that says more about the proxy than about the code.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?

Additional Information

• Fixes the memory_usage regression on basic/tracked-finetune-gemma4-4b-automodel-ac
  (nightly-gb200-finetune-automodel), first seen 2026-08-14.
• Caused by fix(checkpoint): preserve FSDP2 mixed precision during recompute #3513.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.