[yuchenwang3]

PLE backward still allocates whole-sequence FP32 buffers for norm-weight reduction and gate-gradient partials. At 262,144 tokens and width 10,240, each [T, W] FP32 buffer takes 10 GiB.

This accumulates norm-weight and gate-weight gradients in token chunks, with a 64 MiB budget for the three gate partial buffers. It also reuses the consumed norm-output gradient buffer for norm backward. Required input/output gradients and saved activations still scale with the sequence length. Chunked weight reductions change the floating-point summation order, so gradient comparisons use tolerances.

Validation:

• 17 CPU reduction/workspace cases passed on macOS, PyTorch 2.13.0.
• pre-commit run --all-files passed.
• 22 tests passed with no skips on NVIDIA L20C, Python 3.12.3, PyTorch 2.11.0+cu130 / CUDA 13.0. This includes the three added CUDA cases and two existing Torch-reference forward/backward cases, covering in-place norm backward and all seven input gradients. The suite also includes CPU reduction/workspace checks.
• Supplemental FP32/BF16 checks passed against Torch and single-chunk behavior, including BF16 [1024, 10240] inputs. BF16 tolerances were atol=0.002, rtol=0.01; this is tolerance-based gradient agreement, not bitwise equality.
• With a forced 128 KiB workspace budget, the [4096, 512] norm-weight reduction's additional allocated-memory peak fell from 33,559,040 to 90,112 bytes. This measures that helper's workspace, not total model memory or the default-budget saving.
• All six existing opt-in raw-kernel address-boundary checks passed at [262144, 10240]. Those exercise individual kernels, not the complete chunked autograd path or a 256K training step.
• No CP/TP or 256K full-model training run has validated this change yet.

Implementation and tests were generated with OpenAI Codex.

[yuchenwang3]

@hjh0119 Could you review the PLE backward chunking? The GPU run passes all 22 tests, including Torch-reference forward/backward checks. CP/TP and a full 256K training run are still untested.