[Wings-Of-Disaster]

PR type

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

PR information

Fix three sources of unnecessary memory use in the Megatron loss and evaluation path. Each fix is a separate commit.

Base: origin/main at 96dab686.

Area	Problem	Change
selective_log_softmax	The row-by-row loop retains full-vocabulary intermediates until backward.	Process rows in chunks and recompute softmax during backward.
Causal attention mask	The NPU/MindSpeed path builds a dense mask that it later discards. The removal hook also receives a missing config value and returns early.	Default decoder-only causal_lm to causal attention and skip mask creation in collate_fn when the backend supplies it.
Reported logits	Every micro-batch retains detached logits, even when return_logits=False. Fixed-length batches also concatenate them before discarding them.	Retain and collect logits only when requested.

Implementation and compatibility

• Softmax temporaries use a 64 MiB chunk budget, reduced to 32 MiB when entropy is requested. Computation uses float32 for bfloat16 inputs.
• Bfloat16 inputs now produce float32 log-probabilities and entropy. Float32 and float64 output dtypes are unchanged. Existing in-repo callers support this; external callers requiring the old dtype must cast explicitly.
• Mask creation and the fallback removal hook share the same condition. The existing NPU restriction is preserved.
• The collection loop skips None values for both logits and log-probabilities. This also handles failed micro-batches under TWINKLE_FAIL_FAST=0 without passing None to torch.cat.
• Public signatures, defaults, return_logits=True, and the tensor-parallel early return are unchanged. Loss functions still access logits directly.

Mask correctness

For right-padded sequences, the explicit mask and the backend's causal mask differ only at padded query positions. Real tokens cannot attend to later padding, and padded positions are excluded from the loss by labels = -100. Padded outputs must remain finite: multiplying a NaN by a zero loss mask would still produce a NaN.

The reported device check found a maximum output difference of 0.0 at real positions and 1.50 at padded positions. The inspected FlashAttention path uses the causal flag and does not read the explicit mask.

Experiment results

Environment: PyTorch 2.10.0, Transformer Engine 2.9.0, FlashAttention 2.8.2. Vocabulary size: 151,936.

Memory

Saved tensor storage was measured with saved_tensors_hooks, deduplicated by storage pointer, and excludes the input storage. Measurements used 256 rows; the figures below are extrapolated to 16,384 rows. These are retained-storage estimates, not measured accelerator peak memory.

bfloat16 input	Before	After	Save calls at 256 rows
Without entropy	4.64 GiB	0	257 → 1
With entropy	9.27 GiB	0	513 → 1

Float32 without entropy already retained no extra full-vocabulary storage; its benefit is fewer save calls. Input tensors, gradients, and temporary buffers are not included in the zero figures above.

Two additional allocation estimates:

• Unrequested logits: two bfloat16 micro-batches of 16,384 tokens retain 9.27 GiB. Concatenation can raise this to 18.55 GiB when variable_seq_lengths=False. The fix avoids these reporting allocations.
• Dense mask: at batch size 1 and sequence length 16,384, the boolean mask occupies 256 MiB. Creation uses approximately 512 MiB at peak. The fix skips this allocation on the eligible NPU/MindSpeed path.

Numerical accuracy and validation

• On the tested bfloat16 inputs, maximum error against a float64 reference decreased from 3.4e-02 to 7.4e-03.
• Float64 gradcheck passed with and without entropy. Direct gradient comparisons against log_softmax → gather gave maximum differences of 8.941e-08 and 2.235e-08, respectively.
• Processor and utility tests: 74 passed, including 11 added cases.
• Loss, utility, and processor tests: 191 passed, 6 skipped, 3 failed. The same three test_fallback_* failures in test_liger_fused_linear_ce.py occur on unmodified main.
• pre-commit run --all-files: all hooks passed.

Example: why the row-by-row loop retains memory

The previous bfloat16 path computed a full-vocabulary result for each row, then selected only the target token's log-probability:

for row_logits, row_labels in zip(logits, index, strict=True):
    row_logps = F.log_softmax(row_logits, dim=-1)
    # Gather the target token's log-probability from row_logps.

Selecting one value does not release the rest of the row: log_softmax saves its output for backward. At 16,384 rows, all 16,384 outputs remain in memory, equivalent to another full logits tensor:

16,384 × 151,936 × 2 bytes = 4.64 GiB

The replacement recomputes softmax in chunks during backward instead of retaining those outputs.

Example: where the two causal masks differ

Consider right-padded sequences with lengths [4, 2] and a padded length of 4. For the second sequence, positions 0–1 are real tokens and 2–3 are padding.

With True meaning blocked, the explicit mask is T[r,c] = (c > r) OR (c >= 2). The backend's causal mask is C[r,c] = (c > r).

Query row	Token type	Difference
0	Real	None
1	Real	None
2	Padding	T additionally blocks column 2
3	Padding	T additionally blocks columns 2 and 3

The masks are not identical, but real-token outputs are unchanged: each query uses its own mask row, and causal attention prevents real tokens from reading later padding in subsequent layers. Padded positions have labels = -100 and are excluded from the loss.

The finiteness condition still matters. For this example, both masks leave at least one accessible key in every row, avoiding fully masked rows. Loss masking alone would not protect against NaNs because 0 × NaN = NaN.

Example: why the computation and output use float32

The original numerical checks illustrate two different limits:

Case	float16	bfloat16	float32
Flat row, vocabulary 151,936: computed log-probability	-inf	-11.938	-11.931
Peaked row, log-probability near −100.3: maximum error	1.25e-02	2.00e-01	3.05e-06

For a flat row, sum(exp(x - max(x))) = 151,936, exceeding float16's maximum finite value of 65,504 if the sum is stored in float16. The overflow result above describes the tested computation, not every softmax implementation.

Bfloat16 has enough range, but adjacent representable values near −100 are 0.5 apart. Casting the result back to bfloat16 would lose the precision gained during computation. This is why bfloat16 inputs now return float32 log-probabilities and entropy.