[Bug] FLCE Triton Fallback - #1396
Open
buffett0323 wants to merge 2 commits into
Open
Conversation
Contributor
Author
3 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
The CuTe DSL backend's
LigerFusedLinearCrossEntropyFunctioncurrently raises a hardRuntimeErroron capability misses (e.g., non-SM100 hardware, FP32, or unsupported reductions) instead of falling back to Triton. This breaks fused-linear-cross-entropy training across all models when using thecutedslbackend on non-SM100 hardware (reproducible on H100 viaLIGER_KERNEL_IMPL=cutedsl python -m pytest test/convergence/bf16/test_mini_models.py -k "mini_llama3 or mini_qwen2").This PR makes capability misses fall back to the Triton kernel with a one-time warning. Found while validating #1390 against both backends, not specific to any model. The README already sets this expectation for ops with no CuTe DSL kernel, this PR extends it to an op that has one which declines the current hardware.
Details
Dispatch Implementation
LigerFusedLinearCrossEntropyFunctionis now a thin dispatcher rather than atorch.autograd.Function. The original autograd logic is renamed to_NativeFusedLinearCrossEntropyFunction..applyinside an autogradforward()would execute outside the graph and silently drop gradients.__module__assertions remain unchanged.fused_linear_cross_entropy_forwardstill raises on misses; only the dispatcher degrades.Testing Done
Validated across H100 (SM90) and B200 (SM100) to ensure both fallback and native execution paths work properly.
test/cutedslFLCE + convergence (mini_llama3,mini_qwen2,muse_glimmer)test/cutedsl/test_fused_linear_cross_entropy.py(incl.@sm100_required)Test Changes
3 tests now asset Triton fallback instead of a hard raise (
test_flce_falls_back_to_triton_on_non_sm100,..._for_fp32,..._for_reduction_none).test_flce_uses_native_path_when_supported(@sm100_required) to prevent accidental permanent fallback.test_flce_native_forward_still_raises_on_capability_missto check the low-level entry point.test_cutedsl_flce_integration_does_not_raise_on_capability_missto verify the end-to-end integration path via a subprocess.run
make testto ensure correctnessrun
make checkstyleto ensure code stylerun
make test-convergenceto ensure convergence