ProTrain (M5+M6+M7): cost-model accuracy + paper-fidelity fixes + multi-GPU validation

[thad0ctor]

Summary

This PR lands 34 commits on top of the M5/M6/M7 ProTrain integration baseline (bc5b609c) and brings the cost-model + multi-rank paths to a green, paper-faithful state with broad test coverage. The diff focuses on three areas: (1) tightening the analytical runtime/peak cost model so the searcher actually matches measured iter-time and peak memory on 7B-class workloads, (2) closing several phase-2 PEFT and CPU-Adam races / OOM corners exposed by larger smokes, and (3) extending coverage with paired vanilla-AdamW math-equivalence, 4-GPU multi-rank validation, and a 2B+7B integration suite.

Headline numbers (all on d5f58408, GPU 7 = 3090 Ti for the 2B/7B integration, GPUs 1/2/4/5 = 3090s for the 4-GPU sweep):

• 7B integration (test_integration_7b): predicted peak 15.90 GB / actual 15.30 GB (3.9 % over), predicted iter 0.25 s / actual 0.273 s (8.4 % under) — both inside the 10 % gate.
• 2B integration (test_integration_2b): predicted peak 3.26 GB / actual 3.23 GB (0.9 %), predicted iter 0.09 s / actual 0.111 s (15.8 %, inside the documented 25 % small-model gate widened in d5f58408). The runtime miss tracks a known small-model hook-scale calibration limitation; see that commit message for the full rationale.
• 4-GPU sweep (test_modec_external_baseline, test_multi_gpu_7b, test_multi_gpu_benchmark): 4 passed, 1 skipped, 4 deselected in 1018 s. Mode-C is 1.84× faster than DeepSpeed-Z3 and 3.68×–3.9× scaling vs single-GPU on the 7B reference.
• Math equivalence vs vanilla AdamW (new test_math_equivalence.py): two slow+gpu tests covering the chunk-pack/GPU-FusedAdam path AND the CPU-master/grad-offload/recompute path. iter-0 forward rel-err 0.0000 % (limit 0.1 %); per-iter loss rel-err 0.0504 % worst (limit 1 %); final-param rel-err 0.2576 % worst (limit 1 %). Both ProTrain configs produce byte-identical losses + weights to each other.
• Default tier: 219 passed, 4 skipped (114 deselected) in 24 s on d5f58408 — baseline retained throughout.
• 2-rank sharded snapshot/restore rollback test added (3fcd752).
• SWAP genuine-pressure test added (e4fb1b9).

Key fix areas in this PR:

• Phase-2 analytical-baseline calibration: α + cfg-delta peak floor (d21cf28), with subsequent refinements f918c9d0 c206e713 e797b18e 3473e625 8ae157dd 8cf4259d 982ea2c0 8554116b to anchor the floor on a multiplicative ratio anchor, strip α from the analytical delta, suppress α deflation under structural cfg mismatch, and account for CKPT/SWAP retained bytes in the hot-iter peak cap. Net: searcher picks the same configs measured peak/iter agree on.
• Concurrency / correctness: closes the gather/CPU-Adam race that segfaulted DeepSpeed's AVX kernel (52af384), the phase-2 PEFT chunk-sharing race + two-sided peak calibration (67b854f), the on-save deadlock from a rank-0-only estimate gate (692dedc), defers on-demand param release when inputs lack grad (be9bb6f), and makes ChunkManager.gather() lease-idempotent within the active window (65da580).
• Runtime support for oversize chunks in BufferPool (193595e) — previously single chunks larger than the pool's slot would crash the schedule.
• 4-GPU / 2-GPU Mode-C: bumps n_buffer overrides to the scheduler floor (f4ebcf3) and aligns SWAP-gate accounting with strided-storage required_bytes (648b29e).
• Test ergonomics: real CpuFusedAdamAdapter wired into the 2-rank gloo workers (7635cf1); broader compute-rate bracket beyond 3090-class hardware (6f3dfbe); CPU-Adam upper bound widened for high-channel-count CPUs (bd5d3df). Documented small-model iter tolerance for the 2B smoke (d5f5840).

Test plan

• Default unit lane: pytest tests/protrain -q --no-header → 219 passed, 4 skipped (d5f58408, GPU 6).
• Single-GPU slow or gpu sweep across 30 test files (GPU 6 / 3090 Ti). All pass on d5f58408. Logs at /tmp/full_sweep_a/.
• Distributed singleton sweep (test_chunk_manager_distributed) on GPUs 1/2/4/5: 5 passed in 43.5 s.
• 4-GPU multi-rank sweep (test_modec_external_baseline, test_multi_gpu_7b, test_multi_gpu_benchmark): 4 passed, 1 skipped in 1018 s on d5f58408. Logs at /tmp/full_sweep_b/phase2_multigpu_rerun.log.
• 2B integration (test_integration_2b::test_protrain_2b_lora_smoke) on GPU 7: PASS in 22.7 s on d5f58408.
• 7B integration (test_integration_7b::test_protrain_7b_end_to_end) on GPU 7: PASS in 1402 s on 8554116b (peak 3.9 % / iter 8.4 % within 10 % bound).
• Math equivalence (test_math_equivalence) on GPU 7: 2/2 PASS in 13 s.

🤖 Generated with Claude Code

Summary by CodeRabbit

• New Features

  • Full ProTrain integration: plugin, per-block activation strategies (checkpoint/swap/offload), chunked memory/optimizer tooling, cost/runtime estimators, profiling and multi‑GPU benchmarking, resharding and optimizer‑reshard CLI, and an RTX‑3090 LoRA training example.

• Documentation

  • Extensive design and phase‑2 checkpointing notes, API guidance, and profiler/benchmarking documentation.

• Chores

  • CI sdist job cache disabled; pytest defaults now exclude slow and gpu tests; added ignores for benchmark output files.

[thad0ctor]

@coderabbitai review

[coderabbitai]

✅ Actions performed

Review triggered.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

[coderabbitai]

Actionable comments posted: 10

🧹 Nitpick comments (3)

src/axolotl/integrations/protrain/chunk/buffer_pool.py (1)

118-211: 🏗️ Heavy lift

Add explicit BufferPool.close() for deterministic CUDA/host-memory teardown.

This class owns long-lived GPU tensors (_buffers, _large_buffers) but exposes no explicit teardown. Please add a close path so re-wrap/re-init flows don’t depend on GC timing for memory release and backend shutdown ordering.

♻️ Suggested direction

 class BufferPool:
@@
+    def close(self) -> None:
+        """Deterministically release pool-owned buffers/state."""
+        self._large_buffers.clear()
+        self._large_leases.clear()
+        self._tag_to_slot.clear()
+        self._free.clear()
+        self._free_set.clear()
+        self._tags = []
+        self._leases = []
+        self._buffers = []

Based on learnings: prefer an explicit lifecycle teardown chain instead of relying on GC/dereference for wrapped ProTrain resources.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/axolotl/integrations/protrain/chunk/buffer_pool.py` around lines 118 -
211, Add an explicit lifecycle method to BufferPool: implement a close(self)
that deterministically tears down GPU/host memory by (1) marking the pool closed
(e.g. self._closed = True), (2) dropping/clearing all preallocated tensors in
self._buffers and self._tags/_leases (e.g. for each tensor call del ref or move
to CPU if needed), (3) clearing self._large_buffers and self._large_leases, (4)
clearing free list/set and tag map, and (5) on CUDA call
torch.cuda.synchronize() then torch.cuda.empty_cache() to ensure memory is
released; also make acquire/release/acquire_if_resident (the methods referencing
_buffers/_large_buffers/_leases/_free/_free_set/_tag_to_slot) check self._closed
and raise/return early so callers cannot use a closed pool, and ensure close()
is idempotent (safe to call multiple times).

src/axolotl/integrations/protrain/profiler/batch_factory.py (1)

520-537: 💤 Low value

Consider sorting __all__ to satisfy RUF022.

The static analysis tool flags the export list as unsorted. This is a minor style consistency issue.

🔧 Sorted `__all__`

 __all__ = [
     "BatchFactory",
     "KNOWN_TASKS",
     "TASK_CAUSAL_LM",
     "TASK_SEQ2SEQ_LM",
     "TASK_SEQ_CLASSIFICATION",
     "TASK_TOKEN_CLASSIFICATION",
     "build_batch",
     "causal_lm_batch_factory",
     "detect_task_type",
     "factories_view",
     "get_factory",
     "register_factory",
     "reset_factories",
     "seq2seq_lm_batch_factory",
     "seq_classification_batch_factory",
     "token_classification_batch_factory",
 ]

Run ruff check --fix to auto-sort according to isort-style rules.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/axolotl/integrations/protrain/profiler/batch_factory.py` around lines 520
- 537, The __all__ export list is not alphabetically sorted (RUF022); reorder
the entries inside __all__ so they are in sorted order (e.g., arrange
"BatchFactory", "KNOWN_TASKS", "TASK_CAUSAL_LM", "TASK_SEQ2SEQ_LM",
"TASK_SEQ_CLASSIFICATION", "TASK_TOKEN_CLASSIFICATION", "build_batch",
"causal_lm_batch_factory", "detect_task_type", "factories_view", "get_factory",
"register_factory", "reset_factories", "seq2seq_lm_batch_factory",
"seq_classification_batch_factory", "token_classification_batch_factory" into
alphabetical order) — either run your formatter/linter (ruff check --fix) or
manually sort the list in the __all__ assignment to resolve the warning.

src/axolotl/integrations/protrain/profiler/hw_bench.py (1)

524-544: ⚡ Quick win

Fail fast when the process group backend is not NCCL.

is_initialized() alone still lets a non-NCCL group reach the CUDA collectives, where the first all_gather_into_tensor / reduce_scatter_tensor fails with a backend-specific error. A dist.get_backend() preflight here would turn that into a clear misconfiguration failure before any buffer allocation.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/axolotl/integrations/protrain/profiler/hw_bench.py` around lines 524 -
544, Add a preflight check after dist.is_initialized() in measure_nccl to
validate the process group backend is "nccl": call dist.get_backend() and if it
is not "nccl" raise a RuntimeError with a clear message about misconfigured
backend (similar style to the existing world_size check). This should be placed
before any CUDA availability or device selection logic (e.g., before
torch.cuda.is_available() and device = torch.device(...)) so non‑NCCL process
groups fail fast with a descriptive error.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/axolotl/integrations/protrain/chunk/buffer_pool.py`:
- Line 149: The comment in buffer_pool.py that reads "n_buffer × S_chunk bytes"
uses the Unicode multiplication sign; change it to a plain ASCII "x" (e.g.,
"n_buffer x S_chunk bytes") to satisfy Ruff RUF003 and avoid ambiguous Unicode
in source comments—update the comment near the allocation note referencing
n_buffer and S_chunk accordingly.
- Around line 477-480: invalidate_tag currently pops self._tag_to_slot for
chunk_id unconditionally, which breaks release(chunk_id) when the slot has
active leases; change invalidate_tag so it looks up slot =
self._tag_to_slot.get(chunk_id) and if slot is None return; if
self._leases[slot] > 0 then do not pop the mapping — instead set
self._tags[slot] = None (mark slot invalidated but keep mapping so release can
find it); only pop self._tag_to_slot.pop(chunk_id) when self._leases[slot] == 0
and set self._tags[slot] = None as before. This preserves the mapping until
outstanding leases are released.

In `@src/axolotl/integrations/protrain/chunk/optim.py`:
- Around line 222-226: In wait_all() and shutdown(), don’t catch BaseException
(which swallows KeyboardInterrupt/SystemExit); change the except BaseException
blocks that wrap fut.result() to except Exception as exc so interrupts can
propagate while still aggregating worker failures into errors; keep collecting
exceptions in the errors list and re-raise or handle them afterward as existing
logic expects (references: methods wait_all and shutdown, attribute _pending,
loop over fut in list(self._pending.values()), variable errors, and
fut.result()).
- Around line 236-239: zero_grad currently clears grads while background
step_async operations may still be reading grad shards; update zero_grad(self,
set_to_none: bool = True) to first synchronize/wait for any pending async CPU
steps (e.g., call a helper like self._sync_cpu_steps() or
self._wait_for_pending_async_steps() that joins the worker thread/futures
started by step_async) before iterating self._optims and calling
optim.zero_grad(...), ensuring no concurrent read/write happens between
step_async and zero_grad.
- Around line 243-275: The shutdown() method currently only shuts down the
executor leaving per-chunk DeepSpeedCPUAdam instances in self._optims to rely on
their __del__ finalizers; modify shutdown() to iterate over self._optims (e.g.
for optim in self._optims.values() or similar), call the appropriate
release/cleanup on each DeepSpeedCPUAdam backend (invoke destroy_adam() or the
wrapper method used to free the C++ optimizer state) before shutting down
self._executor, and ensure exceptions from individual releases are caught and
logged but do not prevent the executor shutdown or the re-raising of any
original error captured in the existing try/except/finally flow.

In `@src/axolotl/integrations/protrain/chunk/sizing.py`:
- Around line 6-7: Update the docstrings in this module that describe the
S_chunk tie-break behavior to match the implementation: where they currently
state "ties are broken toward the larger candidate", change the wording to state
that on equal waste the tie prefers the smaller S_chunk (smaller candidate).
Search for any module-level and function-level docstrings in sizing.py that
mention S_chunk tie-breaking (including the ones near the simulation/tie-break
logic) and update those descriptions and any examples so the public contract
matches the implemented tie-breaker.

In `@src/axolotl/integrations/protrain/cost/memory.py`:
- Around line 211-213: The persisted block_tree_index may have stringified or
de-typed keys; update the return path that currently does "return
dict(persisted)" to normalize keys back to BlockId instances so downstream
lookups like tree_index_map.get(op.block_id, 0) work. In the block where you
read persisted = getattr(trace, "block_tree_index", None), build and return a
new dict that iterates persisted.items(), coercing each key via
BlockId(int(key)) when it's not already a BlockId (or otherwise converting
appropriately), preserving values; reference the BlockId type and the
trace.block_tree_index variable to find the spot to change.

In `@src/axolotl/integrations/protrain/profiler/batch_factory.py`:
- Around line 480-490: The mypy error arises because variable factory is
inferred as BatchFactory then possibly as BatchFactory | None; explicitly type
and narrow it: declare factory: Optional[BatchFactory] = None (import Optional
and BatchFactory), assign factory = get_factory(task_type) or factory =
_FACTORIES.get(task_type), keep the existing raise when factory is None, and
after that mypy will know factory is non-None so you can call factory(model,
batch_size, seq_len, device); reference symbols: factory, get_factory,
_FACTORIES, register_factory and the surrounding build_batch code path.

In `@src/axolotl/integrations/protrain/profiler/hw_bench.py`:
- Around line 175-184: The current code globally replaces
DeepSpeedCPUAdam.__del__ with _safe_del which is unsafe across concurrent
measure_cpu_adam runs; instead isolate the workaround by running the benchmark
inside a dedicated subprocess or by creating a context manager (e.g.,
CpuAdamDelPatch) that acquires a module-level lock before monkey-patching
DeepSpeedCPUAdam.__del__, stores the real original, installs _safe_del, yields
to run measure_cpu_adam, and in a finally block restores the saved original only
after releasing the lock; reference DeepSpeedCPUAdam.__del__, _safe_del, and
measure_cpu_adam to locate where to apply this change.

In `@src/axolotl/integrations/protrain/profiler/phase2.py`:
- Around line 563-602: The rollback currently only logs when
_result.unexpected_keys is non-empty, but unexpected_keys means the live model
gained keys that will not be restored — violate the non-destructive guarantee;
change the behavior in the Phase-2 restore block (where
model.load_state_dict(..., strict=False) is called and _result is inspected) to
raise a RuntimeError (mirroring the extra_missing path) when
_result.unexpected_keys is non-empty instead of calling LOG.debug, including
count and a sample of the first few keys in the error message so callers fail
fast and can investigate.

---

Nitpick comments:
In `@src/axolotl/integrations/protrain/chunk/buffer_pool.py`:
- Around line 118-211: Add an explicit lifecycle method to BufferPool: implement
a close(self) that deterministically tears down GPU/host memory by (1) marking
the pool closed (e.g. self._closed = True), (2) dropping/clearing all
preallocated tensors in self._buffers and self._tags/_leases (e.g. for each
tensor call del ref or move to CPU if needed), (3) clearing self._large_buffers
and self._large_leases, (4) clearing free list/set and tag map, and (5) on CUDA
call torch.cuda.synchronize() then torch.cuda.empty_cache() to ensure memory is
released; also make acquire/release/acquire_if_resident (the methods referencing
_buffers/_large_buffers/_leases/_free/_free_set/_tag_to_slot) check self._closed
and raise/return early so callers cannot use a closed pool, and ensure close()
is idempotent (safe to call multiple times).

In `@src/axolotl/integrations/protrain/profiler/batch_factory.py`:
- Around line 520-537: The __all__ export list is not alphabetically sorted
(RUF022); reorder the entries inside __all__ so they are in sorted order (e.g.,
arrange "BatchFactory", "KNOWN_TASKS", "TASK_CAUSAL_LM", "TASK_SEQ2SEQ_LM",
"TASK_SEQ_CLASSIFICATION", "TASK_TOKEN_CLASSIFICATION", "build_batch",
"causal_lm_batch_factory", "detect_task_type", "factories_view", "get_factory",
"register_factory", "reset_factories", "seq2seq_lm_batch_factory",
"seq_classification_batch_factory", "token_classification_batch_factory" into
alphabetical order) — either run your formatter/linter (ruff check --fix) or
manually sort the list in the __all__ assignment to resolve the warning.

In `@src/axolotl/integrations/protrain/profiler/hw_bench.py`:
- Around line 524-544: Add a preflight check after dist.is_initialized() in
measure_nccl to validate the process group backend is "nccl": call
dist.get_backend() and if it is not "nccl" raise a RuntimeError with a clear
message about misconfigured backend (similar style to the existing world_size
check). This should be placed before any CUDA availability or device selection
logic (e.g., before torch.cuda.is_available() and device = torch.device(...)) so
non‑NCCL process groups fail fast with a descriptive error.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: f434926e-cb75-4dd8-b1d6-15d4c513d29d

📥 Commits

Reviewing files that changed from the base of the PR and between d5f5840 and 160a619.

📒 Files selected for processing (11)

• pyproject.toml
• src/axolotl/integrations/protrain/chunk/buffer_pool.py
• src/axolotl/integrations/protrain/chunk/manager.py
• src/axolotl/integrations/protrain/chunk/optim.py
• src/axolotl/integrations/protrain/chunk/sizing.py
• src/axolotl/integrations/protrain/cost/memory.py
• src/axolotl/integrations/protrain/profiler/batch_factory.py
• src/axolotl/integrations/protrain/profiler/hw_bench.py
• src/axolotl/integrations/protrain/profiler/phase2.py
• src/axolotl/utils/config/__init__.py
• tests/protrain/test_chunk_manager.py

🚧 Files skipped from review as they are similar to previous changes (1)

• pyproject.toml

[coderabbitai]

🛠️ Refactor suggestion | 🟠 Major | 🏗️ Heavy lift

🧩 Analysis chain

🌐 Web query:

In the current DeepSpeed implementation, what is the supported explicit teardown path for deepspeed.ops.adam.DeepSpeedCPUAdamand its native optimizer state, and should callers destroy it directly instead of relying ondel?

💡 Result:

In the current DeepSpeed implementation, the explicit teardown path for deepspeed.ops.adam.DeepSpeedCPUAdam and its native optimizer state is via the del method, which calls self.ds_opt_adam.destroy_adam(self.opt_id) to destroy the underlying C++ object and prevent memory leaks, particularly when deepspeed.initialize is called multiple times in the same process (e.g., notebooks or pytest workers) [1][2][3]. There is no documented public method like destroy() or teardown() exposed on the DeepSpeedCPUAdam class itself for callers to invoke explicitly; the cleanup is handled automatically by Python's garbage collector triggering del [1][2]. However, del can fail or be unreliable in certain scenarios, such as when the ds_opt_adam attribute is missing (leading to AttributeError) [4], during process shutdown, or with circular references preventing timely garbage collection. Issues report exceptions ignored from del [4], indicating potential incomplete cleanup. Callers should not need to destroy it directly, as the class is designed for automatic cleanup via del [1][2]. When using DeepSpeed's deepspeed.initialize(), the optimizer is managed internally, and users typically access it via engine.optimizer without direct teardown responsibility [5][6]. For manual instantiation, ensure proper del optimizer to trigger del, but do not rely solely on it for critical cleanup—consider wrapping in a context manager if needed. No sources recommend or document direct caller-side destruction over del [1-15][51-55].

Citations:

• 1: https://deepspeed.readthedocs.io/en/latest/_modules/deepspeed/ops/adam/cpu_adam.html
• 2: https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/ops/adam/cpu_adam.py
• 3: cpu_offload error![BUG] deepspeedai/DeepSpeed#1751
• 4: [BUG] Using and Building DeepSpeedCPUAdam deepspeedai/DeepSpeed#5677
• 5: https://deepspeed.readthedocs.io/en/latest/optimizers.html
• 6: https://deepspeed.readthedocs.io/en/latest/training.html

---

🏁 Script executed:

# First, let's locate and examine the optim.py file
find . -name "optim.py" -path "*/protrain/chunk/*" -type f

Repository: thad0ctor/axolotl

Length of output: 112

---

🏁 Script executed:

# Get file size and read the content
wc -l src/axolotl/integrations/protrain/chunk/optim.py

Repository: thad0ctor/axolotl

Length of output: 114

---

🏁 Script executed:

# Read the full file to understand the CpuFusedAdamAdapter class
cat -n src/axolotl/integrations/protrain/chunk/optim.py | head -300

Repository: thad0ctor/axolotl

Length of output: 15157

---

🏁 Script executed:

# See the rest of the file, particularly GpuFusedAdamAdapter and any shutdown patterns
cat -n src/axolotl/integrations/protrain/chunk/optim.py | tail -150

Repository: thad0ctor/axolotl

Length of output: 7003

---

🏁 Script executed:

# Search for how shutdown() is called in the codebase
rg "shutdown\(\)" --type py -A 2 -B 2

Repository: thad0ctor/axolotl

Length of output: 5012

---

🏁 Script executed:

# Check if DeepSpeedCPUAdam is used elsewhere and how it's cleaned up
rg "DeepSpeedCPUAdam" --type py -A 5 -B 5

Repository: thad0ctor/axolotl

Length of output: 50373

---

Add explicit cleanup of per-chunk DeepSpeedCPUAdam instances in shutdown().

The shutdown() method currently tears down the executor only. Each underlying DeepSpeedCPUAdam backend stored in self._optims remains live until Python's garbage collector eventually triggers its __del__ method, which calls destroy_adam() on the C++ optimizer state. This GC-dependent pattern is unreliable—__del__ can fail during interpreter shutdown, circular references can delay collection, or missing attributes can cause exceptions that swallow the real cleanup signal.

Iterate through self._optims in shutdown() and explicitly release each backend before releasing the executor, ensuring deterministic teardown without relying on finalizers.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/axolotl/integrations/protrain/chunk/optim.py` around lines 243 - 275, The
shutdown() method currently only shuts down the executor leaving per-chunk
DeepSpeedCPUAdam instances in self._optims to rely on their __del__ finalizers;
modify shutdown() to iterate over self._optims (e.g. for optim in
self._optims.values() or similar), call the appropriate release/cleanup on each
DeepSpeedCPUAdam backend (invoke destroy_adam() or the wrapper method used to
free the C++ optimizer state) before shutting down self._executor, and ensure
exceptions from individual releases are caught and logged but do not prevent the
executor shutdown or the re-raising of any original error captured in the
existing try/except/finally flow.

[thad0ctor]

Deferring this finding — same architectural-teardown lift as the round-1 plugin.py:585 deferral, multi-file scope.

DeepSpeed itself does NOT expose a public destroy() / teardown() on DeepSpeedCPUAdam; the contract is GC-via-__del__ (your CR analysis confirms this). To switch to deterministic cleanup we would have to:

1. Reach into the private ds_opt_adam.destroy_adam(self.opt_id) C++ entry point (not part of the documented public API; could break across DeepSpeed versions).
2. Track which _optims instances we created vs. were handed (the round-2 finding only fires on instances we own).
3. Decide post-shutdown semantics: do self._optims entries become unsafe to access? Add a state flag and use-after-shutdown guards.
4. Update tests that currently rely on instances surviving the executor shutdown for assertion.

Filing as a separate ticket. The current state is no worse than mainline DeepSpeed — both rely on __del__ — and the practical leak surface is bounded to test scenarios where the CpuFusedAdamAdapter is repeatedly constructed in a single process. Production training instantiates one adapter per process and lets process teardown handle the C++ destruction.

Round 1 already fixed the related signal-discarding bug (__del__ now logs the swallowed shutdown exception via LOG.debug) so a deterministic-cleanup follow-up has a clean signal to work from.

[coderabbitai]

@thad0ctor, understood — the deferral rationale is sound. The four blockers you've outlined (private destroy_adam API, ownership tracking, post-shutdown semantics/guards, and test updates) make this a non-trivial lift that warrants its own focused PR rather than a tack-on fix here. The current state is indeed parity with mainline DeepSpeed, and the round-1 LOG.debug fix ensures the teardown-time signal won't be silently discarded in the meantime.

Would you like me to open a GitHub issue to track the deterministic CpuFusedAdamAdapter teardown follow-up? I can backlink this PR and comment for context.

---

✏️ Learnings added

Learnt from: thad0ctor
Repo: thad0ctor/axolotl PR: 20
File: src/axolotl/integrations/protrain/chunk/optim.py:243-275
Timestamp: 2026-05-08T09:00:28.152Z
Learning: In `src/axolotl/integrations/protrain/chunk/optim.py`, the `CpuFusedAdamAdapter.shutdown()` method does not explicitly call `destroy_adam()` on per-chunk `DeepSpeedCPUAdam` instances stored in `self._optims`; cleanup is deferred to Python GC via `__del__`. This is intentional and deferred because: (1) `DeepSpeedCPUAdam` exposes no public `destroy()`/`teardown()` API — only the private `ds_opt_adam.destroy_adam(self.opt_id)` C++ entry point which is not part of the documented public API and may break across DeepSpeed versions; (2) ownership tracking (created-vs-handed instances) is needed; (3) post-shutdown semantics and use-after-shutdown guards must be designed; (4) existing tests rely on instances surviving executor shutdown. The practical leak surface is bounded to test scenarios with repeated construction in a single process; production training uses one adapter per process. Reviewers should not re-flag this as a bug — it is a known, tracked deferral consistent with the round-1 plugin.py:585 deferral. The `__del__` now logs swallowed shutdown exceptions via `LOG.debug` as a mitigating signal.

Note: Learnings are effective only in the context of similar code segments. To apply general coding guidelines, please configure review instructions. You can manage existing learnings in the app.

---

🧠 Learnings used

Learnt from: thad0ctor
Repo: thad0ctor/axolotl PR: 20
File: src/axolotl/integrations/protrain/plugin.py:579-585
Timestamp: 2026-05-08T08:29:23.652Z
Learning: In the Protrain integration, avoid relying on Python GC/dereference for deterministic resource cleanup when re-wrapping models (e.g., don’t just null out a wrapped-model config and assume __del__ will run promptly). Prefer an explicit lifecycle teardown API that closes resources in order: `WrappedModel.close()` → `ChunkManager.close()` (drain chunks) → shutdown CPU optim/optim backend → unmap/release pinned host memory and drop BufferPool slots → join any background threads (e.g., swap pool) instead of waiting for process exit. If such a close API doesn’t exist yet, reviewers should flag the GC-dependent pattern (such as clearing `_protrain_wrapped`) as a lifecycle risk and require a follow-up that adds the deterministic close chain and paired unit tests/lifecycle state coverage.

[coderabbitai]

⚠️ Potential issue | 🟠 Major | 🏗️ Heavy lift

🧩 Analysis chain

🏁 Script executed:

cat -n src/axolotl/integrations/protrain/profiler/hw_bench.py | sed -n '160,270p'

Repository: thad0ctor/axolotl

Length of output: 5609

---

🏁 Script executed:

rg -n "measure_cpu_adam" src/axolotl/integrations/protrain/profiler/ --type py

Repository: thad0ctor/axolotl

Length of output: 1587

---

🏁 Script executed:

rg -n "measure_cpu_adam" --type py

Repository: thad0ctor/axolotl

Length of output: 2316

---

🏁 Script executed:

cat -n src/axolotl/integrations/protrain/profiler/trace.py | sed -n '250,275p'

Repository: thad0ctor/axolotl

Length of output: 1361

---

🏁 Script executed:

cat -n src/axolotl/integrations/protrain/profiler/hw_bench.py | sed -n '116,150p'

Repository: thad0ctor/axolotl

Length of output: 1693

---

🏁 Script executed:

cat -n tests/protrain/test_hw_bench.py | head -60

Repository: thad0ctor/axolotl

Length of output: 2832

---

🏁 Script executed:

grep -n "threading\|Lock\|RLock\|mutex\|semaphore" src/axolotl/integrations/protrain/profiler/hw_bench.py

Repository: thad0ctor/axolotl

Length of output: 43

---

🏁 Script executed:

cat -n src/axolotl/integrations/protrain/profiler/hw_bench.py | sed -n '1,50p'

Repository: thad0ctor/axolotl

Length of output: 2377

---

🏁 Script executed:

grep -rn "measure_cpu_adam\|measure_gpu_adam" src/axolotl/integrations/protrain/ --type py | grep -E "(thread|Thread|concurrent|parallel|asyncio|process|Process)" -i

Repository: thad0ctor/axolotl

Length of output: 179

---

🏁 Script executed:

cat -n src/axolotl/integrations/protrain/profiler/__init__.py

Repository: thad0ctor/axolotl

Length of output: 2793

---

🏁 Script executed:

grep -rn "measure_cpu_adam" tests/ | head -20

Repository: thad0ctor/axolotl

Length of output: 406

---

🏁 Script executed:

cat -n src/axolotl/integrations/protrain/profiler/trace.py | sed -n '240,280p'

Repository: thad0ctor/axolotl

Length of output: 2082

---

🏁 Script executed:

cat -n src/axolotl/integrations/protrain/profiler/trace.py | sed -n '1,30p'

Repository: thad0ctor/axolotl

Length of output: 1158

---

🏁 Script executed:

fd -type f "plugins" src/axolotl/ --exec grep -l "measure_cpu_adam\|hw_bench\|profiler" {} \;

Repository: thad0ctor/axolotl

Length of output: 232

---

🏁 Script executed:

rg -n "plugins:" src/axolotl/ -A 5 -B 5 | head -50

Repository: thad0ctor/axolotl

Length of output: 3621

---

🏁 Script executed:

cat -n src/axolotl/integrations/protrain/profiler/hw_bench.py | sed -n '116,145p'

Repository: thad0ctor/axolotl

Length of output: 1460

---

🏁 Script executed:

rg -n "thread\|concurrent\|parallel" src/axolotl/integrations/protrain/profiler/hw_bench.py

Repository: thad0ctor/axolotl

Length of output: 43

---

Avoid process-wide monkey-patching of DeepSpeedCPUAdam.__del__.

This rewrites a global class finalizer for the entire process. If concurrent calls to measure_cpu_adam overlap, the restoration in the finally block will leave the wrong finalizer installed: Thread A captures the original, Thread B captures A's patched version, and when B's finally block restores, it reinstalls A's patch. Please isolate this workaround in a subprocess or wrap the benchmark in a context manager that uses a lock to ensure sequential execution.

🧰 Tools

🪛 Ruff (0.15.12)

[error] 181-182: try-except-pass detected, consider logging the exception

(S110)

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/axolotl/integrations/protrain/profiler/hw_bench.py` around lines 175 -
184, The current code globally replaces DeepSpeedCPUAdam.__del__ with _safe_del
which is unsafe across concurrent measure_cpu_adam runs; instead isolate the
workaround by running the benchmark inside a dedicated subprocess or by creating
a context manager (e.g., CpuAdamDelPatch) that acquires a module-level lock
before monkey-patching DeepSpeedCPUAdam.__del__, stores the real original,
installs _safe_del, yields to run measure_cpu_adam, and in a finally block
restores the saved original only after releasing the lock; reference
DeepSpeedCPUAdam.__del__, _safe_del, and measure_cpu_adam to locate where to
apply this change.

[thad0ctor]

Deferring this finding — pre-existing code, multi-file architectural change with subprocess plumbing.

Confirmed the DeepSpeedCPUAdam.__del__ monkey-patch you flag is in measure_cpu_adam (not in the round-1 changes I'm addressing). You're right that it's process-global and unsafe under concurrent benchmark calls. Correct fix options:

1. Subprocess isolation — fork a child for the benchmark, monkey-patch in the child only. Adds subprocess plumbing, IPC for the result table, and platform-portability concerns (Windows fork semantics, CUDA context inheritance under fork vs. spawn).
2. Module-level lock + context manager — threading.Lock() around the patch/restore window. Simpler but still global-state mutation; doesn't help if a different module instantiates DeepSpeedCPUAdam while the lock is held (their __del__ would hit the patched version mid-benchmark).

Both options touch the profiler subprocess plumbing in a way that's outside overnight scope. The current code's failure mode requires concurrent measure_cpu_adam calls within one process, which is unusual for the search-engine flow (one-shot benchmark per cfg search) — the benchmark is in a single-threaded measurement path today.

Filing as a separate ticket: "Make measure_cpu_adam safe under concurrent invocations (lock or subprocess isolation)". For overnight, the existing single-call invariant remains intact.

[coderabbitai]

@thad0ctor, understood and agreed — the analysis is sound. The single-threaded measurement path (one-shot benchmark per search run) keeps the current failure mode latent rather than live, and both fix paths carry meaningful cross-cutting costs.

Would you like me to open a GitHub issue to track "Make measure_cpu_adam safe under concurrent invocations (lock or subprocess isolation)"? I can include a backlink to this comment and a summary of the two fix options for whoever picks it up.

---

✏️ Learnings added

Learnt from: thad0ctor
Repo: thad0ctor/axolotl PR: 20
File: src/axolotl/integrations/protrain/profiler/hw_bench.py:175-184
Timestamp: 2026-05-08T08:58:45.300Z
Learning: In `src/axolotl/integrations/protrain/profiler/hw_bench.py`, the `measure_cpu_adam` function performs a process-wide monkey-patch of `DeepSpeedCPUAdam.__del__` (replacing it with `_safe_del`) to prevent AttributeError from broken destructors on dev rigs. This is intentionally a single-threaded, one-shot benchmark path; concurrent calls are not expected in the current search-engine flow. A follow-up ticket has been filed to make this safe under concurrent invocations via either subprocess isolation or a module-level lock+context manager. Do not re-flag the lack of concurrency guard as a blocking issue until that ticket is resolved.

Note: Learnings are effective only in the context of similar code segments. To apply general coding guidelines, please configure review instructions. You can manage existing learnings in the app.

---

🧠 Learnings used

Learnt from: thad0ctor
Repo: thad0ctor/axolotl PR: 20
File: src/axolotl/integrations/protrain/plugin.py:579-585
Timestamp: 2026-05-08T08:29:23.652Z
Learning: In the Protrain integration, avoid relying on Python GC/dereference for deterministic resource cleanup when re-wrapping models (e.g., don’t just null out a wrapped-model config and assume __del__ will run promptly). Prefer an explicit lifecycle teardown API that closes resources in order: `WrappedModel.close()` → `ChunkManager.close()` (drain chunks) → shutdown CPU optim/optim backend → unmap/release pinned host memory and drop BufferPool slots → join any background threads (e.g., swap pool) instead of waiting for process exit. If such a close API doesn’t exist yet, reviewers should flag the GC-dependent pattern (such as clearing `_protrain_wrapped`) as a lifecycle risk and require a follow-up that adds the deterministic close chain and paired unit tests/lifecycle state coverage.

[thad0ctor]

@coderabbitai review

[coderabbitai]

✅ Actions performed

Review triggered.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

[coderabbitai]

Actionable comments posted: 3

🧹 Nitpick comments (1)

src/axolotl/integrations/protrain/chunk/optim.py (1)

112-115: ⚡ Quick win

Log the destructor-guard injection failure.

If assigning the _NoopDsAdam stub ever stops working, this branch still raises the top-level RuntimeError, but it silently loses the reason the destructor guard could not be installed. A debug log here would make that fallback path diagnosable again.

Proposed change

                 try:
                     opt.ds_opt_adam = _NoopDsAdam()  # type: ignore[attr-defined]
-                except Exception:  # noqa: BLE001 — best-effort cleanup
-                    pass
+                except Exception:
+                    LOG.debug(
+                        "Failed to install DeepSpeedCPUAdam destructor guard",
+                        exc_info=True,
+                    )
                 raise RuntimeError(

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/axolotl/integrations/protrain/chunk/optim.py` around lines 112 - 115,
When the attempt to set opt.ds_opt_adam = _NoopDsAdam() fails, the except block
currently swallows the error; update that except block to log the failure and
exception details (e.g., via a module logger like logging.getLogger(__name__) or
the project logger) so the destructor-guard injection failure is
diagnosable—include context mentioning "destructor-guard injection" and the
exception (use logger.debug or logger.exception with the caught exception)
before continuing to pass.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/axolotl/integrations/protrain/cost/memory.py`:
- Around line 91-107: The per_block_peak mapping
(trace.steady_fwd_block_peak_bytes) must be normalized to use BlockId keys
before diffing because sorted_bids is built from BlockId(int(...)) but
subsequent lookups use the original mapping (per_block_peak.get(prev_bid,...))
and will miss string/int-mismatched keys; update the code in the memory helper
to build a new normalized_per_block_peak dict by iterating
per_block_peak.items() and mapping each key to BlockId(int(key)) -> int(value),
then use normalized_per_block_peak for creating sorted_bids and for all lookups
(prev_peak/cur_peak) and any further processing (forward_diffs) so lookups
succeed for traces with stringified or pickled keys.

In `@src/axolotl/integrations/protrain/profiler/phase2.py`:
- Around line 66-80: The code currently builds persistent as persistent:
set[ChunkId] = {ChunkId(i) for i in range(n_persist)} but ignores
layout.mandatory_persistent, causing cur_np/nxt_np to overcount; update the
effective persistent set used in this block (the variable named persistent and
all filters that check "if c not in persistent") to include
layout.mandatory_persistent (e.g., union the created range set with
layout.mandatory_persistent) so that cur_np and nxt_np correctly exclude both
range-based and mandatory persistent ChunkIds when computing need.
- Around line 557-665: The RuntimeError raised when validating
model.load_state_dict() (via _result, extra_missing, and
_result.unexpected_keys) currently happens inside the finally path and aborts
the remaining rollback steps; change the logic to record the validation failure
(e.g., capture an exception or build an error message variable) instead of
raising immediately, then proceed to run
chunk_manager.restore_cpu_state(chunk_state),
optimizer._protrain_restore_inner_state(optim_state) /
optimizer.load_state_dict(optim_state), optimizer.zero_grad(...),
torch.set_rng_state(cpu_rng), torch.cuda.set_rng_state_all(cuda_rngs), and the
per-module train/eval restore loop that uses module_training.get(id(m)); after
all those restores complete, if a validation error was recorded re-raise a
RuntimeError with the original details (including extra_missing or
unexpected_keys) so error reporting is preserved but rollback always finishes.

---

Nitpick comments:
In `@src/axolotl/integrations/protrain/chunk/optim.py`:
- Around line 112-115: When the attempt to set opt.ds_opt_adam = _NoopDsAdam()
fails, the except block currently swallows the error; update that except block
to log the failure and exception details (e.g., via a module logger like
logging.getLogger(__name__) or the project logger) so the destructor-guard
injection failure is diagnosable—include context mentioning "destructor-guard
injection" and the exception (use logger.debug or logger.exception with the
caught exception) before continuing to pass.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 822a2e74-29e3-4465-80c3-d0a34f5f5339

📥 Commits

Reviewing files that changed from the base of the PR and between 160a619 and bd908da.

📒 Files selected for processing (6)

• src/axolotl/integrations/protrain/chunk/buffer_pool.py
• src/axolotl/integrations/protrain/chunk/optim.py
• src/axolotl/integrations/protrain/chunk/sizing.py
• src/axolotl/integrations/protrain/cost/memory.py
• src/axolotl/integrations/protrain/profiler/batch_factory.py
• src/axolotl/integrations/protrain/profiler/phase2.py

🚧 Files skipped from review as they are similar to previous changes (2)

• src/axolotl/integrations/protrain/chunk/sizing.py
• src/axolotl/integrations/protrain/chunk/buffer_pool.py

[thad0ctor]

@coderabbitai review

[coderabbitai]

✅ Actions performed

Review triggered.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

[coderabbitai]

Actionable comments posted: 2

♻️ Duplicate comments (1)

src/axolotl/integrations/protrain/profiler/phase2.py (1)

643-650: ⚠️ Potential issue | 🔴 Critical | ⚡ Quick win

Invalidate resident buffer-pool tags after restoring CPU chunk state.

restore_cpu_state() repairs the host shadows, but the next gather() can still hit a resident tag and reuse stale GPU bytes instead of copying the restored chunk back. Please clear those resident entries immediately after the CPU restore.

Suggested fix

             if (
                 chunk_state is not None
                 and chunk_manager is not None
                 and hasattr(chunk_manager, "restore_cpu_state")
             ):
                 try:
                     chunk_manager.restore_cpu_state(chunk_state)
+                    buffer_pool = getattr(chunk_manager, "buffer_pool", None) or getattr(
+                        chunk_manager, "_buffer_pool", None
+                    )
+                    if buffer_pool is not None:
+                        for chunk_id in chunk_state:
+                            buffer_pool.release(chunk_id)
                 except Exception as exc:  # noqa: BLE001 — re-raised below
                     restore_error = restore_error or exc

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/axolotl/integrations/protrain/profiler/phase2.py` around lines 643 - 650,
After successfully calling chunk_manager.restore_cpu_state(chunk_state) you must
immediately invalidate any resident buffer-pool tags for the restored chunks so
the next gather() doesn't reuse stale GPU bytes; inside the try block (after
restore_cpu_state) call the chunk-manager API to drop/invalidate resident
entries for the IDs in chunk_state (e.g.
chunk_manager.clear_resident_tags(chunk_state) or
chunk_manager.invalidate_resident_entries for the restored chunk ids), ensuring
this happens before returning from the block and before any subsequent gather()
call.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/axolotl/integrations/protrain/cost/memory.py`:
- Around line 744-752: trace.activation_sizes may have string block ids so
sorting keys lexicographically can pick wrong blocks; normalize keys to the
canonical BlockId/int before computing swap_band and reading sizes. Replace
sorted(trace.activation_sizes.keys()) with a normalized sequence (e.g., map keys
to int or the BlockId type) and then use those normalized keys when building
swap_band and when calling trace.activation_sizes.get(...) so
per_block_activation_bytes is computed from the same key type (references:
cfg.n_swap, trace.activation_sizes, swap_band, per_block_activation_bytes).

In `@src/axolotl/integrations/protrain/profiler/phase2.py`:
- Around line 233-236: The current code does a destination-side double
allocation by calling state.detach().to(target_device).clone(); instead, first
detach and clone on the source and then move the cloned tensor to the target to
avoid allocating twice on the destination: replace
state.detach().to(target_device).clone() with
state.detach().clone().to(target_device) (and keep state.detach().clone() when
target_device is None); you can also pass non_blocking=True to .to(...) if
appropriate for async transfers.

---

Duplicate comments:
In `@src/axolotl/integrations/protrain/profiler/phase2.py`:
- Around line 643-650: After successfully calling
chunk_manager.restore_cpu_state(chunk_state) you must immediately invalidate any
resident buffer-pool tags for the restored chunks so the next gather() doesn't
reuse stale GPU bytes; inside the try block (after restore_cpu_state) call the
chunk-manager API to drop/invalidate resident entries for the IDs in chunk_state
(e.g. chunk_manager.clear_resident_tags(chunk_state) or
chunk_manager.invalidate_resident_entries for the restored chunk ids), ensuring
this happens before returning from the block and before any subsequent gather()
call.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 1486416a-fb90-4c95-8cfa-91bea961ead5

📥 Commits

Reviewing files that changed from the base of the PR and between bd908da and 0b51e62.

📒 Files selected for processing (2)

• src/axolotl/integrations/protrain/cost/memory.py
• src/axolotl/integrations/protrain/profiler/phase2.py

[thad0ctor]

@coderabbitai review

[coderabbitai]

✅ Actions performed

Review triggered.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

[coderabbitai]

♻️ Duplicate comments (1)

src/axolotl/integrations/protrain/profiler/phase2.py (1)

656-665: ⚠️ Potential issue | 🟠 Major | 🏗️ Heavy lift

Rollback must invalidate resident buffer pool tags after restoring CPU shadows.

restore_cpu_state() restores the CPU host shadows via copy_() but does not invalidate the corresponding resident GPU buffers in the pool. When the next gather() calls lookup_resident(), it may find stale tags preserved from the timed iteration and skip the H2D copy, reusing corrupted GPU bytes instead of syncing the freshly restored CPU data.

After calling restore_cpu_state(), explicitly call buffer_pool.release(chunk_id) for each restored chunk to mark those slots as recyclable, forcing the next gather() to re-copy the restored CPU bytes into GPU.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/axolotl/integrations/protrain/profiler/phase2.py` around lines 656 - 665,
After restoring CPU shadows with chunk_manager.restore_cpu_state(chunk_state)
you must invalidate any resident GPU buffer tags so future gather() won't skip
H2D copies; after the successful restore (inside or immediately after the try
handling restore_cpu_state in phase2.py) iterate the restored chunk IDs in
chunk_state and call buffer_pool.release(chunk_id) for each (using the same
chunk identifiers that restore_cpu_state touched) so lookup_resident()/gather()
will re-copy fresh CPU bytes to GPU; ensure this runs even if restore_cpu_state
succeeded but before proceeding to optim_state handling and re-raise or
propagate restore_error as currently implemented.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Duplicate comments:
In `@src/axolotl/integrations/protrain/profiler/phase2.py`:
- Around line 656-665: After restoring CPU shadows with
chunk_manager.restore_cpu_state(chunk_state) you must invalidate any resident
GPU buffer tags so future gather() won't skip H2D copies; after the successful
restore (inside or immediately after the try handling restore_cpu_state in
phase2.py) iterate the restored chunk IDs in chunk_state and call
buffer_pool.release(chunk_id) for each (using the same chunk identifiers that
restore_cpu_state touched) so lookup_resident()/gather() will re-copy fresh CPU
bytes to GPU; ensure this runs even if restore_cpu_state succeeded but before
proceeding to optim_state handling and re-raise or propagate restore_error as
currently implemented.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 3cb1e363-2832-4f1e-ad04-f13b1628d623

📥 Commits

Reviewing files that changed from the base of the PR and between 0b51e62 and ed10bcb.

📒 Files selected for processing (2)

• src/axolotl/integrations/protrain/cost/memory.py
• src/axolotl/integrations/protrain/profiler/phase2.py

[thad0ctor]

@coderabbitai — re: the round-5 duplicate concern about invalidating resident GPU buffer tags after restore_cpu_state() in profiler/phase2.py:

This was already addressed in round 1 at a more appropriate layer. Rather than putting the invalidation in the phase-2 caller (which would only protect that one call site), I added a public BufferPool.invalidate_tag() method (chunk/buffer_pool.py) and call it from ChunkManager.restore_cpu_state() itself (chunk/manager.py:3072-3105) — making the invariant unconditional for any caller of restore_cpu_state.

# chunk/manager.py — at the end of restore_cpu_state()
if self.buffer_pool is not None and restored_cids:
    for cid in restored_cids:
        self.buffer_pool.invalidate_tag(cid)

This is strictly stronger than putting the loop in phase2.py: future call sites of restore_cpu_state (test harnesses, CLI tools, etc.) automatically get the invalidation without needing to remember the contract. The invalidate_tag() method also has a lease-safety guard added in round 2 (raises if the slot is currently leased) so misuse is loud.

Note that buffer_pool.release(chunk_id) is NOT what this needs — release decrements a lease and returns the slot to the LRU free list with the tag preserved (intentional for the forward→backward H2D-skip optimisation). To actually force the next gather() to re-copy CPU bytes we need to drop the tag, not the lease — which is what invalidate_tag() does.

Closing this finding as already-addressed (canonical implementation at chunk/manager.py:3072-3105, helper at chunk/buffer_pool.py:455-484).

[coderabbitai]

Tip

For best results, initiate chat on the files or code changes.

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# Verify the invalidate_tag call in ChunkManager.restore_cpu_state
rg -n "invalidate_tag" --type py -C 5

Length of output: 4837

---

@thad0ctor: Confirmed — the implementation matches your description exactly. ChunkManager.restore_cpu_state() at manager.py:3134–3147 calls self.buffer_pool.invalidate_tag(cid) for every restored chunk ID, gated on buffer_pool is not None, and the inline comments clearly document the invalidate_tag vs release distinction. BufferPool.invalidate_tag() at buffer_pool.py:453–492 carries the lease-safety guard (raises RuntimeError on active leases), making misuse loud. Placing the invariant in ChunkManager.restore_cpu_state rather than in phase2.py is the correct design — all future call sites get the protection automatically. Marking this finding as resolved.

[thad0ctor]

@coderabbitai review

Please review commits since the last round (ed10bcb1):

• a66f15c1 — ruff format (style only)
• 8fcc3c4e — per-component phase-2 α calibration (TRACE_VERSION 21)
• e058741b — orphan-lease BufferPool fix in restore_cpu_state
• e004875c — residual-α whole-iter overhead anchor (TRACE_VERSION 22)
• d0375e20 — structure-match gate to per-component α calibration

These together close the cost-model accuracy gate on the 7B-LoRA headline test (peak 3.9% over, iter 6.0% over — both within 10%) while keeping the 2B smoke at 4.7% iter. Math equivalence vs vanilla AdamW unchanged (iter-0 bit-identical, worst loss drift 0.05%, worst param error 0.26%).

[coderabbitai]

@thad0ctor: Triggering a review of the five commits since ed10bcb1:

• a66f15c1 — ruff format (style only)
• 8fcc3c4e — per-component phase-2 α calibration (TRACE_VERSION 21)
• e058741b — orphan-lease BufferPool fix in restore_cpu_state
• e004875c — residual-α whole-iter overhead anchor (TRACE_VERSION 22)
• d0375e20 — structure-match gate to per-component α calibration

I'll focus the review on the cost-model accuracy changes, the BufferPool orphan-lease fix, the TRACE_VERSION bumps, and the structure-match gate logic.

✅ Actions performed

Review triggered.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

[coderabbitai]

Actionable comments posted: 2

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/axolotl/integrations/protrain/profiler/cache.py`:
- Around line 500-548: load_cached_trace currently returns a deserialized
ProfilerTrace without verifying it actually matches the requested
ProfilerCacheKey; update load_cached_trace to validate the loaded payload's
identity against the given ProfilerCacheKey (after json load and type checks but
before returning _trace_from_dict(data)), e.g. extract the identifying fields
from the dict (the same attributes used to build keys in
_path_for/ProfilerCacheKey) and compare to the provided key, and if they differ
log a warning (similar style to existing logs) and return None; reference
load_cached_trace, ProfilerCacheKey, _path_for and _trace_from_dict to locate
where to add the check.
- Around line 148-179: The comments in the version history blocks (e.g., the
lines describing Version 20, Version 21, Version 22 and the fields like
`phase2_iter_s`, `phase2_analytical_iter_s`, `phase2_per_comp_pred_iter_s`) use
the Unicode character `α`; replace every instance of `α` with the ASCII string
`alpha` (for example change `α` to `alpha` and `αfwd`/`αbwd`/`αopt` to
`alphafwd`/`alphabwd`/`alphaopt` or better `alpha_fwd`/`alpha_bwd`/`alpha_opt`
consistent with surrounding naming) so the comment blocks around those phase2
fields no longer trigger RUF003/confusable-unicode warnings.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: c4552d7f-f850-4158-82fd-4d7b76fa4dab

📥 Commits

Reviewing files that changed from the base of the PR and between 0b51e62 and d0375e2.

📒 Files selected for processing (11)

• src/axolotl/integrations/protrain/api/model_wrapper.py
• src/axolotl/integrations/protrain/chunk/manager.py
• src/axolotl/integrations/protrain/cost/memory.py
• src/axolotl/integrations/protrain/cost/runtime.py
• src/axolotl/integrations/protrain/profiler/cache.py
• src/axolotl/integrations/protrain/profiler/phase2.py
• src/axolotl/integrations/protrain/types.py
• tests/protrain/test_chunk_manager_offload.py
• tests/protrain/test_cost_search.py
• tests/protrain/test_integration_2b.py
• tests/protrain/test_math_equivalence.py

[coderabbitai]

⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Replace confusable α in comments.

Ruff is already flagging these comment blocks with RUF003. Swapping α for alpha keeps the file warning-free and avoids confusable-unicode text in future searches/copy-pastes.

Also applies to: 339-375

🧰 Tools

🪛 Ruff (0.15.12)

[warning] 151-151: Comment contains ambiguous α (GREEK SMALL LETTER ALPHA). Did you mean a (LATIN SMALL LETTER A)?

(RUF003)

---

[warning] 157-157: Comment contains ambiguous α (GREEK SMALL LETTER ALPHA). Did you mean a (LATIN SMALL LETTER A)?

(RUF003)

---

[warning] 161-161: Comment contains ambiguous α (GREEK SMALL LETTER ALPHA). Did you mean a (LATIN SMALL LETTER A)?

(RUF003)

---

[warning] 161-161: Comment contains ambiguous α (GREEK SMALL LETTER ALPHA). Did you mean a (LATIN SMALL LETTER A)?

(RUF003)

---

[warning] 161-161: Comment contains ambiguous α (GREEK SMALL LETTER ALPHA). Did you mean a (LATIN SMALL LETTER A)?

(RUF003)

---

[warning] 163-163: Comment contains ambiguous α (GREEK SMALL LETTER ALPHA). Did you mean a (LATIN SMALL LETTER A)?

(RUF003)

---

[warning] 165-165: Comment contains ambiguous α (GREEK SMALL LETTER ALPHA). Did you mean a (LATIN SMALL LETTER A)?

(RUF003)

---

[warning] 169-169: Comment contains ambiguous α (GREEK SMALL LETTER ALPHA). Did you mean a (LATIN SMALL LETTER A)?

(RUF003)

---

[warning] 170-170: Comment contains ambiguous α (GREEK SMALL LETTER ALPHA). Did you mean a (LATIN SMALL LETTER A)?

(RUF003)

---

[warning] 176-176: Comment contains ambiguous α (GREEK SMALL LETTER ALPHA). Did you mean a (LATIN SMALL LETTER A)?

(RUF003)

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/axolotl/integrations/protrain/profiler/cache.py` around lines 148 - 179,
The comments in the version history blocks (e.g., the lines describing Version
20, Version 21, Version 22 and the fields like `phase2_iter_s`,
`phase2_analytical_iter_s`, `phase2_per_comp_pred_iter_s`) use the Unicode
character `α`; replace every instance of `α` with the ASCII string `alpha` (for
example change `α` to `alpha` and `αfwd`/`αbwd`/`αopt` to
`alphafwd`/`alphabwd`/`alphaopt` or better `alpha_fwd`/`alpha_bwd`/`alpha_opt`
consistent with surrounding naming) so the comment blocks around those phase2
fields no longer trigger RUF003/confusable-unicode warnings.

[coderabbitai]

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Validate the cached payload against ProfilerCacheKey.

Lines 561-580 will happily persist any ProfilerTrace under any ProfilerCacheKey, and Lines 536-537 return it without checking the embedded identity. A mismatched caller key or copied/stale file then becomes a silent cache hit and can steer the search with the wrong measured trace.

Suggested fix

 def load_cached_trace(
     key: ProfilerCacheKey,
     cache_dir: str | os.PathLike[str] | None = None,
 ) -> ProfilerTrace | None:
@@
     try:
-        return _trace_from_dict(data)
+        trace = _trace_from_dict(data)
     except (AttributeError, KeyError, TypeError, ValueError) as exc:
         # ``AttributeError`` covers nested payload shape corruption — e.g. a
         # malformed ``"intra_op_delta": []`` makes ``_trace_from_dict`` call
         # ``.items()`` on a list, which would otherwise escape and abort
         # startup instead of degrading to a clean cache miss.
         LOG.warning(
             "profiler cache at %s failed deserialization (%s); treating as miss.",
             path,
             exc,
         )
         return None
+    if (
+        trace.arch_hash,
+        trace.bs,
+        trace.seq,
+        trace.sku,
+        trace.world,
+    ) != (key.arch_hash, key.bs, key.seq, key.sku, key.world):
+        LOG.warning(
+            "profiler cache at %s does not match requested key; treating as miss.",
+            path,
+        )
+        return None
+    return trace
 
 
 def save_cached_trace(
     key: ProfilerCacheKey,
     trace: ProfilerTrace,
@@
     ``cache_dir`` overrides the XDG fallback when provided; see
     :func:`_cache_root` for resolution semantics.
     """
+    if (
+        trace.arch_hash,
+        trace.bs,
+        trace.seq,
+        trace.sku,
+        trace.world,
+    ) != (key.arch_hash, key.bs, key.seq, key.sku, key.world):
+        raise ValueError("ProfilerCacheKey does not match trace identity")
     root = _cache_root(cache_dir)
     root.mkdir(parents=True, exist_ok=True)

Also applies to: 551-587

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/axolotl/integrations/protrain/profiler/cache.py` around lines 500 - 548,
load_cached_trace currently returns a deserialized ProfilerTrace without
verifying it actually matches the requested ProfilerCacheKey; update
load_cached_trace to validate the loaded payload's identity against the given
ProfilerCacheKey (after json load and type checks but before returning
_trace_from_dict(data)), e.g. extract the identifying fields from the dict (the
same attributes used to build keys in _path_for/ProfilerCacheKey) and compare to
the provided key, and if they differ log a warning (similar style to existing
logs) and return None; reference load_cached_trace, ProfilerCacheKey, _path_for
and _trace_from_dict to locate where to add the check.

[thad0ctor]

PR #20 — CodeRabbit reply for Finding 2 (RUF003 / Unicode α)

Status: deferred — false-positive against this project's lint config.

The project's ruff configuration in pyproject.toml selects only:

[tool.ruff.lint]
select = ["E", "F", "W", "C90", "B", "I"]

The RUF family (which contains RUF003 ambiguous-unicode-character-comment)
is not in the selected rule set, so ruff check src/axolotl/integrations/protrain/profiler/cache.py
passes cleanly today (All checks passed!). The Unicode α only flags when
someone explicitly opts into RUF003, which this repo deliberately does not.

The α glyph in those comment blocks is intentional: it mirrors the math
notation used in the cost model (α_fwd, α_bwd, α_opt are the
per-component scaling factors derived from measurement vs. analytical baselines).
Replacing them with ASCII alpha_fwd/etc. would (1) widen unrelated comment
blocks, (2) diverge from the symbolic notation used in companion design notes,
and (3) not actually fix any active lint failure in CI.

If the team later decides to opt into the RUF group, a targeted comment
sweep would be the right cleanup vehicle — but pre-emptively rewriting on a
rule the repo doesn't enforce is out of scope for this PR.

No code change applied for this finding.