Repository navigation
[dist] refactor: inject the async-offload host buffer pool - #1175
Conversation
OffloadManager owned its PinnedBufferPool, which tied a process-wide memory budget to a per-schedule object. A manager has to stay per schedule -- block_idx means "position within this schedule", so two sets numbering from 0 would prefetch and wait on each other's tensors -- but the pool does not: its buffers are keyed by layout alone and are fungible across schedules. Injecting the pool lets a composed model apply offload per module without multiplying activation_offload_host_cache_limit_gb by the module count. Single-model callers keep passing host_cache_limit_bytes and are unaffected. Also folds in the f-string conversions the pyupgrade hook applies to this file. Co-authored-by: Cursor <cursoragent@cursor.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Advanced Run ID: 📒 Files selected for processing (4)
🚧 Files skipped from review as they are similar to previous changes (3)
Included review availability: Your plan provides up to 8 included reviews per hour; 7 remain after this review. 📝 WalkthroughWalkthroughAsync activation offload now supports shared ChangesActivation offload pool lifecycle
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant apply_async_activation_offload
participant PinnedBufferPool
participant async_offload_modules
participant OffloadManager
Caller->>apply_async_activation_offload: configure offload
apply_async_activation_offload->>PinnedBufferPool: resolve or create shared pool
apply_async_activation_offload->>async_offload_modules: pass resolved pool
async_offload_modules->>OffloadManager: construct manager with pool
OffloadManager->>PinnedBufferPool: reuse pool for offload schedules
Merge Risk: ⚪ Minimal · up to Activation offload can now use an injected shared pinned-buffer pool or construct an independent pool per call, with conflicting options rejected and cache limits applied at the pool level. The documented behavior is covered by unit tests, and no current merge-blocking risk is identified. 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Full details: Docstring CoverageExplanation Docstring coverage is 43.75% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 16 functions across 3 files. (1 skipped: 1 unsupported.)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
@CodeRabbit |
|
Your plan includes PR reviews subject to rate limits. Reviews are available now. |
…-Seed#1175) Co-authored-by: Cursor <cursoragent@cursor.com>
What does this PR do?
Split out of #1090:
OffloadManagerno longer owns itsPinnedBufferPool. The manager stays per offload schedule (block_idxis “position within this schedule”), while the host pool is process-wide and keyed only by tensor layout, so a composed model can apply offload per module without multiplyingactivation_offload_host_cache_limit_gbby the module count.Single-model callers keep passing
host_cache_limit_bytesand are unchanged. A caller that applies offload more than once (omni per-module, DPO reference, …) passes one sharedhost_buffer_pool.Checklist Before Starting
[{modules}] {type}: {description}format (enforced by check_pr_title.yml)agent,ci,ckpt,config,data,dist,docker,docs,logging,lora,misc,model,omni,optim,ops,parallel,perf,release,task,trainerfeat,fix,refactor,chore,test[BREAKING]— e.g.[BREAKING][parallel, model] feat: dynamic batchingTest
Extended the existing CI-enumerated file
tests/distributed/test_async_offload_unit.py(gpu_unit_tests.ymlalready lists it). Added cases for: a shared pool reused across twoapply_async_activation_offloadcalls, building a pool fromhost_cache_limit_bytes, rejecting a non-poolOffloadManagerargument, and rejecting limit+pool together.CPU path:
pytest tests/distributed/test_async_offload_unit.py -k "not accelerator_gradient_parity and not fsdp2_gc"— 27 passed, 2 skipped (device tests). The GPU integration tests in that file are unchanged and still run in GPU CI.API and Usage Example
Single-model (unchanged):
Several schedules sharing one host budget:
host_cache_limit_bytesandhost_buffer_poolare mutually exclusive. YAMLactivation_offload_host_cache_limit_gbis still the process-wide idle-cache cap; only the help text clarifies that.Design & Code Changes
OffloadManagernow takes aPinnedBufferPoolinstead of creating one from a byte limit.async_offload_modules(..., *, host_buffer_pool=...)is keyword-only for the pool.apply_async_activation_offloadbuilds a pool when given a limit, or uses a caller-supplied pool. Passing both raisesValueError.OffloadConfighelp: the GB limit is counted once per process, not per offloaded module.reset_async_activation_offloadstill walks one root; a multi-root caller must reset each root.Checklist Before Submitting
tasks/training scripts were moved or renamed: updateddocs/examples and verifiedpython3 scripts/ci/check_doc_task_paths.pypasses (also enforced by the Check doc task paths CI workflow)tests/distributed/test_async_offload_unit.py, already listed ingpu_unit_tests.yml).Summary by CodeRabbit
New Features
Documentation