fix(dev): support current torchao and torchrl releases - #2982
Open
n-dlms wants to merge 1 commit into
Open
Conversation
torchao >= 0.17 moved the NF4 API (NF4Tensor, to_nf4, linear_nf4, implements) from torchao.quantization / quantize_.workflows.nf4 to torchao.dtypes.nf4tensor, breaking `import torchtune` against the latest torchao, which the CI workflows install unpinned. torchrl >= 0.13 removed SyncDataCollector, WeightUpdateReceiverBase, WeightUpdateSenderBase, and LLMEnv, breaking the async RL data path against released torchrl. - Import the NF4 symbols with a fallback to torchao.dtypes.nf4tensor - Add dev/rl/utils/_torchrl_compat.py with import-compatible stand-ins for the removed torchrl collector classes - Pin the async_rl extra to torchrl 0.13.3 / tensordict 0.13.0 instead of stale git commits (the extra's own 'update to a stable release soon' TODO) - Drop Python 3.9 from the RL CI matrix (torchrl >= 0.13 requires 3.10+) Verified against torchao 0.17.0 + torchrl 0.13.3 + tensordict 0.13.0: tests/torchtune/dev/rl passes 43, skips 1 (GPU-gated).
meta-cla
Bot
added
the
CLA Signed
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.
Summary
torchtune currently fails to import and its CI is red against current releases of its two unmanaged dependencies:
NF4Tensor,to_nf4,linear_nf4,implements) fromtorchao.quantization/torchao.quantization.quantize_.workflows.nf4.nf4_tensortotorchao.dtypes.nf4tensor. All CI workflows (unit_test,gpu_test,regression_test,rl_test,build_docs) install torchao unpinned, soimport torchtunebreaks on the latest release.SyncDataCollector,WeightUpdateReceiverBase,WeightUpdateSenderBase, andLLMEnv, which the async RL data path subclasses/imports. Theasync_rlextra still pins two stale git commits with the note "We will update this to a stable release soon!".Changes
try/exceptfallback totorchao.dtypes.nf4tensorin the seven modules that use them (keeps older torchao 0.7-0.16 working).torchtune/dev/rl/utils/_torchrl_compat.py: imports the torchrl collector classes when available, otherwise provides import-compatible stand-ins so the RL data path can be imported and unit-tested against released torchrl.pyproject.toml: theasync_rlextra now pinstorchrl==0.13.3/tensordict==0.13.0instead of the stale git commits..github/workflows/rl_test.yaml: drop Python 3.9 from the matrix (torchrl >= 0.13 requires Python >= 3.10).Notes
SyncLLMCollector,VLLMHFWeightUpdateReceiver) still requires torchrl < 0.13 with the original collector classes; the stand-ins intentionally raiseNotImplementedErrorthere rather than pretending to work.pytest tests/torchtune/dev/rlpasses 43, skips 1 (GPU-gated), including when combined with fix(grpo): replay buffer contract corruptions advantages (#2943) #2981's data-path tests.