Repository navigation
Conversation
Documentation build overview
68 files changed ·
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces an opt-in wall-clock span measurement utility (offload_spans) to profile the performance of the GPU offload and onload operations (sleep and wake_up) in colocate mode, along with corresponding unit tests and pipeline integration. The review feedback highlights two key improvement opportunities: first, resolving a calculation bug where nested spans are double-counted in the total accounted time, and second, refactoring the unit tests to use monkeypatch.setattr instead of importlib.reload to prevent potential state leakage and simplify the test setup.
| # The recorded spans are disjoint segments of one call, so their sum is what | ||
| # the breakdown actually accounts for; the remainder is call overhead. Summing | ||
| # the rounded entries keeps `accounted_s` equal to the table next to it. | ||
| accounted = round(sum(entry["total_s"] for entry in payload.values()), 6) |
There was a problem hiding this comment.
The current calculation of accounted_s sums all recorded spans. However, the module explicitly supports nested spans (as stated in the docstring and tested in test_parent_includes_children). If nested spans are used, summing all values will double-count the nested child spans, leading to an incorrect accounted_s (which can exceed total_s and result in a negative unaccounted_s).
To fix this, we should only sum the top-level (root) spans by filtering out any span that has a parent span in the payload.
# Only sum top-level spans to avoid double-counting nested spans
accounted = round(
sum(
entry["total_s"]
for name, entry in payload.items()
if not any(name.startswith(parent + ".") for parent in payload if parent != name)
),
6,
)| @pytest.fixture(autouse=True) | ||
| def _restore_default_module(): | ||
| """Every case reloads the module; leave it as the rest of the suite found it.""" | ||
| yield | ||
| importlib.reload(offload_spans) | ||
|
|
||
|
|
||
| def _module(monkeypatch, *, enabled: bool): | ||
| if enabled: | ||
| monkeypatch.setenv("VIME_OFFLOAD_SPANS", "1") | ||
| else: | ||
| monkeypatch.delenv("VIME_OFFLOAD_SPANS", raising=False) | ||
| return importlib.reload(offload_spans) |
There was a problem hiding this comment.
Reloading modules in tests using importlib.reload is generally an anti-pattern that can lead to subtle bugs and state leakage, especially when other modules import functions from it using from ... import ... (as those imports won't be updated by the reload).
Instead of reloading the module to toggle _ENABLED, we can use monkeypatch.setattr to dynamically set _ENABLED on the imported module. Since monkeypatch automatically restores the original value after each test, we can also completely remove the _restore_default_module fixture.
| @pytest.fixture(autouse=True) | |
| def _restore_default_module(): | |
| """Every case reloads the module; leave it as the rest of the suite found it.""" | |
| yield | |
| importlib.reload(offload_spans) | |
| def _module(monkeypatch, *, enabled: bool): | |
| if enabled: | |
| monkeypatch.setenv("VIME_OFFLOAD_SPANS", "1") | |
| else: | |
| monkeypatch.delenv("VIME_OFFLOAD_SPANS", raising=False) | |
| return importlib.reload(offload_spans) | |
| def _module(monkeypatch, *, enabled: bool): | |
| monkeypatch.setattr(offload_spans, "_ENABLED", enabled) | |
| return offload_spans |
7d98d35 to
4bfd159
Compare
In colocate mode every step hands the GPUs back and forth between Megatron and vLLM, and Actor.sleep() / Actor.wake_up() sit on the critical path doing it. Their @Timer totals are the largest single block of a small-model short-response step, but they say nothing about where inside the two calls the time goes: sleep() alone is clear_memory(clear_host_memory=True) + destroy_process_groups() + torch_memory_saver.pause() + two print_memory() calls, and those have very different fixes. Add the breakdown, off by default behind VIME_OFFLOAD_SPANS=1: * vime/observability/offload_spans.py accumulates perf_counter() deltas per span name and emits one line per sleep()/wake_up() call with the per-segment totals, the accounted share and the residual call overhead. Disabled it records nothing and emits nothing. * Actor.sleep()/wake_up() wrap each step. The segments are disjoint, so their sum is directly comparable to the call's own total. Measured on 1 GPU colocate, Qwen3-0.6B, rollout_batch_size 8 x n_samples 2, cap 192, 8 rollouts: sleep median 11158 ms of which torch_memory_saver.pause() is 9681 ms (86.8%); wake_up median 3772 ms of which resume() is 3173 ms (84.1%). Process-group lifetime is not the cost (destroy 914 ms + reload 8 ms), and print_memory is free (~5 ms across both calls). The transfer moves ~18.7 GB at 1.67 GB/s out and 4.56 GB/s back over a Gen4 x16 link, so the headroom is in the transfer itself rather than in the code around it. Signed-off-by: 0z5a <0z5a@users.noreply.github.com> Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
4bfd159 to
01dc1b7
Compare
Signed-off-by: 0z5a <0z5a@users.noreply.github.com>
Signed-off-by: 0z5a <0z5a@users.noreply.github.com>
Signed-off-by: 0z5a <0z5a@users.noreply.github.com>
Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
Incremental review: Files changed.
Base:
0z5a-branchat791f101e20db377b28e15f12a00863baaea2605f.Purpose
Colocated training currently reports only aggregate
Actor.sleep()andActor.wake_up()times. This change exposes which operations consume that time: memory-saver pause/resume, memory cleanup, process-group teardown/reload, model switching and diagnostic calls.VIME_OFFLOAD_SPANS=1enables wall-time spans and one summary per call. The default path skips the extra clock reads and emits nothing. Nested spans retain inclusive child/parent totals, whileaccounted_scounts root intervals once. Training, tensor transfers and returned values retain their existing behavior.The submitted diff contains two production files. Test harnesses, logs and their added CI registration were removed; their measured results are retained below.
Validation
Measured results
Default-disabled spans, Qwen3-0.6B, two prompts × two samples, 16 response tokens maximum. Driver round time includes rollout, training, offload and weight publication.
Generated response lengths vary between arms. These timings do not establish a causal training speedup from the accounting change. The CPU microbenchmark observes approximately 1.1 μs per disabled context-manager invocation; enabled nested accounting costs about 2% more than its previous implementation.
With spans enabled in a separate diagnostic round:
sleepwake_upA separate existing
--no-offload-trainconfiguration comparison observed a 4.345× median ratio. That result is not caused by the span patch. Multi-rank coverage remains outstanding.