feat: Optimize memory footprint of long-context training via fused kernel and chunking

[terminator123]

What does this PR do ?

Introduces a fused CrossEntropy kernel and output chunking strategy to reduce the peak memory consumption of logits during long-context training.

Technical Details

This PR addresses the high VRAM usage bottleneck in large-scale training by targeting the logits tensor memory footprint.

• Fused Kernel: Utilizes the Liger-Kernel's fused CrossEntropy implementation to reduce intermediate memory overhead.

• Output Chunking: Implements an output chunking mechanism where the model's output is processed in blocks.

• Memory-Specific Optimization: The peak memory is reduced by a factor proportional to the number of chunks (1/N ). The more chunks the output is divided into, the lower the peak memory .

[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.

[github-actions]

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

1. Add the oncall reviewer (optional reviewer)
2. Add required review teams based on your changes

See the contribution guide for more details.

[Phlip79]

We are in the process of developing this same feature: #2206. @Jianbing-D can you please take a look at this?

[Jianbing-D]

Hi @terminator123,

We have similar feature already merged to dev branch. #2256

And regarding your PR, here are some questions:

1. Is there any measurement numbers regarding your feature? like latency of fwd pass and bwd pass, as well as storage. Like what we did here: [Dev] Feature: linear cross entropy fusion #2256
2. Your kernels are written with OAI triton, but that library failed to achieve good performance on Blackwell GPUs. If you could provide any perf numbers, that would be great for us to determine whether your kernels are good enough.
3. Your feature seems not support reduction=none. Please correct me if I understand it wrong. If this feature doesn't support reduction=none, how could the users handle token masking and padding, where invalid tokens shall have zero as grad, whilst valid tokens shall have valid grad values.