[NOT4REVIEW] Make RL model compilable: ops rewrite + per-sample backward#2394
Draft
Lucaskabela wants to merge 3 commits intopytorch:mainfrom
Draft
[NOT4REVIEW] Make RL model compilable: ops rewrite + per-sample backward#2394Lucaskabela wants to merge 3 commits intopytorch:mainfrom
Lucaskabela wants to merge 3 commits intopytorch:mainfrom
Conversation
meta-cla
bot
added
the
CLA Signed
Lucaskabela
force-pushed
the
lucaskabela/rl-compilable
branch
2 times, most recently
from
57ffef9 to
e10c3b7
Compare
Rewrites batch-invariant ops as torch.library.custom_op (rms_norm, silu_and_mul, flash_attn) so they are opaque to Dynamo/AOT autograd. Adds aten dispatch overrides for matmul/linear backward to use vLLM's deterministic kernels. Refactors compute_policy_gradient_loss_vllm to use per-sample gradient accumulation: each sample's forward is immediately followed by backward, keeping only one set of activations in memory at a time. This is a prerequisite for torch.compile since the compiled graph processes one sample at a time with fixed-shape inputs. Changes: - batch_invariant_backward.py: custom ops rewrite - models/attention.py: custom op for flash_attn - simple_rl.py: per-sample backward, loss_scale param, timing metrics - trainer.py: move zero_grad before loss, remove loss.backward()
Lucaskabela
force-pushed
the
lucaskabela/rl-compilable
branch
from
e10c3b7 to
575a33d
Compare
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
1 participant
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.
Rewrites batch-invariant ops as torch.library.custom_op (rms_norm, silu_and_mul, flash_attn) so they are opaque to Dynamo/AOT autograd. Adds aten dispatch overrides for matmul/linear backward to use vLLM's deterministic kernels.
Refactors compute_policy_gradient_loss_vllm to use per-sample gradient accumulation: each sample's forward is immediately followed by backward, keeping only one set of activations in memory at a time. This is a prerequisite for torch.compile since the compiled graph processes one sample at a time with fixed-shape inputs.
Changes:
Authored with Claude