[DRAFT] Wire through RL compilation: compile_utils, graph_utils, config#5
Draft
Lucaskabela wants to merge 4 commits intolucaskabela/rl-compilablefrom
Draft
[DRAFT] Wire through RL compilation: compile_utils, graph_utils, config#5Lucaskabela wants to merge 4 commits intolucaskabela/rl-compilablefrom
Lucaskabela wants to merge 4 commits intolucaskabela/rl-compilablefrom
Conversation
Lucaskabela
force-pushed
the
lucaskabela/rl-compilable
branch
from
ee75c66 to
57ffef9
Compare
Lucaskabela
force-pushed
the
lucaskabela/rl-compile-wiring
branch
from
86b8622 to
6e0185e
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
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()
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()
Adds compilation support to the RL training loop: - graph_utils.py: return gm from export_joint, add validate_dtensor param, use gm.named_parameters(remove_duplicate=False) to handle tied weights - compile_utils.py: new file using joint_graph_builder with validate_dtensor=False, RLCompiledModule wrapper with input padding - simple_rl.py: add compilation config vars and compile_rl_model call - trainer.py: add compile params, compile before DDP - simple_rl_multiprocess.py: add compilation config and fw+bw timing
Lucaskabela
force-pushed
the
lucaskabela/rl-compile-wiring
branch
from
6e0185e to
962d80f
Compare
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.
Adds compilation support to the RL training loop: