Fix three memory defects in the Megatron path (selective_log_softmax, causal mask, discarded logits)

[Wings-Of-Disaster]

PR type

• Bug Fix
• New Feature
• Document Updates
• More Models or Datasets Support

PR information

Three independent defects in the Megatron training path. All three cost device
memory at long sequence lengths; the first one also costs accuracy. Reproduced
on c839a4e with a Qwen3-4B GRPO configuration (max_train_len=16384,
micro_batch_size=1, mini_batch_size=2, mixed_precision=bf16,
vocab 151936).

1. selective_log_softmax pins one full-vocab copy per row

src/twinkle/utils/torch_utils.py

The per-row loop is commented loop to reduce peak mem consumption. That holds
under no_grad, but not under autograd: every row's log_softmax output (and,
with return_entropy=True, its exp) is handed to save_for_backward, so it
stays resident until the backward pass finishes. The loop therefore pins exactly
what it was written to avoid, one full-vocab activation per row.

Replaced with a torch.autograd.Function that chunks by a byte budget, computes
in float32, and saves only the original-dtype chunk plus its labels, recomputing
the softmax in backward (the standard activation-checkpointing trade).

Side effect, also a fix: the float32 compute removes the bfloat16 quantization
of the log-probability itself. The old bfloat16 branch returned
F.log_softmax(row) in bfloat16, whose ULP is 0.5 at the magnitudes a peaked
row produces.

Behaviour change: log-probs and entropies are returned in float32 for a
bfloat16 input (previously bfloat16). These are per-token tensors, so the cost
is N * 4 bytes and the values entering the loss are no longer quantized.

2. attention_mask_type is always None, so the causal-mask hook is dead code

src/twinkle/model/megatron/megatron.py, src/twinkle/processor/base.py

MegatronModel.forward_backward passed
getattr(unwrapped_model.config, 'attention_mask_type', None) to the processor.
Neither megatron.core.transformer.transformer_config.TransformerConfig nor
mcore_bridge.ModelConfig declares that attribute (AttnMaskType is only an
enum), so the value was None for every model, and
InputProcessor.drop_causal_4d_mask returned early on its
attention_mask_type != 'causal' guard. The dense mask was always built and
always passed down.

Second defect: that hook is pipeline step 7 while collate_fn is step 4, so even
a working hook ran after _create_4d_attention_mask had already allocated
[B, 1, S, S].

The type now defaults to 'causal' for causal_lm, and collate_fn decides up
front, so _create_4d_attention_mask is never called on that path. The hook is
kept as a fallback for callers that build the mask elsewhere. The gate stays
NPU-only (Platform.device_prefix() == 'npu'), i.e. scoped to the backend that
rebuilds its own compressed causal mask, so CUDA behaviour is unchanged.

3. Full-vocab logits are retained, concatenated, then discarded

src/twinkle/model/megatron/megatron.py

MegatronStrategy.reduce_loss puts logits.detach() into every microbatch's
report dict unconditionally; forward_backward appends all of them, torch.cats
the set into a second full-size buffer, and then throws it away one line later
because return_logits defaults to False. The detached references also defeat
the scheduler's own release of each microbatch's output_tensor, so the whole
set is pinned across the backward pass, where memory is tightest. This hits
forward_only too, i.e. the frozen reference model.

TransformersModel.forward has guarded the same thing with
if not (return_logits or loss_require_logits) since 6921342; the Megatron
backend never got the equivalent (git log -S loss_require_logits on
src/twinkle/model/megatron/megatron.py returns nothing). This adds it.

The report dicts are now collected with .get(...) is not None as well:
reduce_loss reports None for both logits and logps when a microbatch's
data processing failed under TWINKLE_FAIL_FAST=0, which the previous
if 'logps' in loss_dict test let through into torch.cat.

Experiment results

1. selective_log_softmax, measured with torch.autograd.graph.saved_tensors_hooks

N=16384 rows, V=151936. "1x" = one full-vocab activation in the input dtype
(4.6 GiB in bf16, 9.3 GiB in fp32). "extra" excludes the caller's own logits
storage, i.e. it is memory this function alone keeps alive until backward ends.

config	extra pinned for backward	save_for_backward calls
before, bf16	1.00x (4.6 GiB)	49152
before, bf16 + entropy	2.00x (9.3 GiB)	98304
before, fp32 + entropy	1.00x (9.3 GiB)	81922
before, fp32	0.00x	32770
after, all four	0.00x	298

Accuracy against a float64 row-wise reference (N=512, V=151936):

dtype	max abs err, logps	max abs err, grads
bf16, before	7.7e-2	1.6e-2 (with entropy)
bf16, after	3.0e-2	7.2e-3 (with entropy)
fp32, before / after	2.0e-6 / 2.1e-6	2.0e-4 / 7.9e-5 (with entropy)

2. Dense 4D mask

_create_4d_attention_mask at B=1, S=16384 on CPU: 256 MiB mask, +499 MiB peak
RSS from the ones / tril / ~ temporaries, ~70 ms per micro-batch. After
the fix the call never happens on the causal NPU path.

3. Discarded logits

2 microbatches x [1, 16384, 151936] bf16: 9.3 GiB of detached per-microbatch
logits pinned past their backward, plus a 9.3 GiB torch.cat buffer, discarded
by if not return_logits: logits = None. 18.6 GiB peak, ~14 GiB of it pure
waste.

Tests

New CPU-only tests, each verified to fail on c839a4e and pass on this branch:

• tests/utils/test_utils.py::TestSelectiveLogSoftmax
  • test_backward_saves_no_extra_full_vocab_activation - fails for bf16 (both
    entropy settings) and fp32 + entropy. fp32 without entropy passes on both:
    that path already pinned nothing extra, and the test says so rather than
    overclaiming.
  • test_scales_with_chunks_not_with_rows - fails on main (1536 saves vs 2 at
    N=512, V=512).
  • test_bfloat16_does_not_quantize_large_negative_logps - fails on main.
• tests/processor/test_processor.py::TestMegatronCausalMask
  • test_causal_npu_omits_dense_mask - fails on main (dense mask allocated).
  • two control tests asserting the mask is still built for non-causal input and
    for non-NPU backends.

Defect 3 has no CPU-testable unit test: post_loss_function is a closure inside
forward_backward and needs an initialised Megatron parallel state. Its evidence
is the code path plus the transformers/megatron asymmetry above.

pytest tests/utils tests/processor tests/loss tests/model tests/advantage tests/metric: 281 passed, 6 skipped. The 5 failures in tests/loss and
tests/model are pre-existing and environmental (they allocate on a CUDA device
that is already full) and fail identically against unmodified c839a4e.