fix: clear_cache_every_n_steps variable name - #1109
Conversation
WalkthroughUpdated both dtensor policy worker training loops to use dtensor_cfg.clear_cache_every_n_steps instead of dtensor_cfg.empty_cache_every_n_steps for periodic CUDA cache clearing. Logic for warnings and torch.cuda.empty_cache() calls remains unchanged. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant Trainer
participant Config as dtensor_cfg
participant CUDA as torch.cuda
Trainer->>Config: Read clear_cache_every_n_steps
alt clear_cache_every_n_steps is set
note right of Trainer: Warning logged about periodic cache clearing
loop Every N steps
Trainer->>CUDA: empty_cache()
CUDA-->>Trainer: cache cleared
end
else not set / zero
note right of Trainer: No periodic cache clearing
end
Estimated code review effortðŊ 2 (Simple) | âąïļ ~10 minutes Poem
Tip ðŪ Agentic pre-merge checks are now available in preview!Pro plan users can now enable pre-merge checks in their settings to enforce checklists before merging PRs.
Example: reviews:
pre_merge_checks:
custom_checks:
- name: "Undocumented Breaking Changes"
mode: "warning"
instructions: |
Pass/fail criteria: All breaking changes to public APIs, CLI flags, environment variables, configuration keys, database schemas, or HTTP/GraphQL endpoints must be documented in the "Breaking Change" section of the PR description and in CHANGELOG.md. Exclude purely internal or private changes (e.g., code not exported from package entry points or explicitly marked as internal).Please share your feedback with us on this Discord post. âĻ Finishing Touches
𧊠Generate unit tests
Comment |
âđïļ File Consistency CheckCheck based on commit: 675b39e (PR #1109 from â DTensor Policy Worker Synchronization CheckBoth DTensor policy worker files were modified in this PR:
Please ensure that the changes are consistent between both files where applicable. This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning. |
675b39e to
ecaf55f
Compare
âđïļ File Consistency CheckCheck based on commit: ecaf55f (PR #1109 from â DTensor Policy Worker Synchronization CheckBoth DTensor policy worker files were modified in this PR:
Please ensure that the changes are consistent between both files where applicable. This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning. |
bxyu-nvidia
changed the title
clear_cache_every_n_steps variable nameclear_cache_every_n_steps variable name
Signed-off-by: Brian Yu <bxyu@nvidia.com>
Signed-off-by: Brian Yu <bxyu@nvidia.com>
âđïļ File Consistency CheckCheck based on commit: 59bfb55 (PR #1109 from â DTensor Policy Worker Synchronization CheckBoth DTensor policy worker files were modified in this PR:
Please ensure that the changes are consistent between both files where applicable. This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning. |
59bfb55 to
7662ab0
Compare
âđïļ File Consistency CheckCheck based on commit: 7662ab0 (PR #1109 from â DTensor Policy Worker Synchronization CheckBoth DTensor policy worker files were modified in this PR:
Please ensure that the changes are consistent between both files where applicable. This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and canât be posted inline due to platform limitations.
â ïļ Outside diff range comments (2)
nemo_rl/models/policy/dtensor_policy_worker_v2.py (1)
577-585: Make the key change backward-compatible and validate input.Switching to "clear_cache_every_n_steps" will silently disable cache clearing for any jobs still using the old "empty_cache_every_n_steps" key, risking more OOMs. Read the new key first, then fall back to the old key with a deprecation warning, and validate the value (int > 0).
Apply:
- empty_cache_steps = self.cfg.get("dtensor_cfg", {}).get( - "clear_cache_every_n_steps" - ) - if empty_cache_steps: - warnings.warn( - f"Emptying cache every {empty_cache_steps} microbatches, doing so unnnecessarily would incur a large performance overhead." - ) + cfg_dt = self.cfg.get("dtensor_cfg", {}) + empty_cache_steps = cfg_dt.get("clear_cache_every_n_steps") + if empty_cache_steps is None: + empty_cache_steps = cfg_dt.get("empty_cache_every_n_steps") + if empty_cache_steps is not None: + warnings.warn( + "dtensor_cfg.empty_cache_every_n_steps is deprecated; use dtensor_cfg.clear_cache_every_n_steps", + DeprecationWarning, + ) + # normalize and validate + if empty_cache_steps is not None: + try: + empty_cache_steps = int(empty_cache_steps) + except (TypeError, ValueError): + warnings.warn( + "dtensor_cfg.clear_cache_every_n_steps must be an int; disabling CUDA cache clearing.", + RuntimeWarning, + ) + empty_cache_steps = None + if empty_cache_steps is not None and empty_cache_steps <= 0: + empty_cache_steps = None + if empty_cache_steps and not getattr(self, "_warned_cache_clear_overhead", False): + warnings.warn( + f"Emptying cache every {empty_cache_steps} microbatches; unnecessary use can incur large performance overhead." + ) + self._warned_cache_clear_overhead = Truenemo_rl/models/policy/dtensor_policy_worker.py (1)
633-640: Preserve compatibility with existing configs and validate the value.Mirroring v2: support both config keys (new preferred: clear_cache_every_n_steps; old deprecated: empty_cache_every_n_steps) to prevent regressions and OOMs if older configs are still in circulation. Add basic type/positivity checks and warn once per process.
- empty_cache_steps = self.cfg.get("dtensor_cfg", {}).get( - "clear_cache_every_n_steps" - ) - if empty_cache_steps: - warnings.warn( - f"Emptying cache every {empty_cache_steps} microbatches, doing so unnnecessarily would incur a large performance overhead." - ) + cfg_dt = self.cfg.get("dtensor_cfg", {}) + empty_cache_steps = cfg_dt.get("clear_cache_every_n_steps") + if empty_cache_steps is None: + empty_cache_steps = cfg_dt.get("empty_cache_every_n_steps") + if empty_cache_steps is not None: + warnings.warn( + "dtensor_cfg.empty_cache_every_n_steps is deprecated; use dtensor_cfg.clear_cache_every_n_steps", + DeprecationWarning, + ) + if empty_cache_steps is not None: + try: + empty_cache_steps = int(empty_cache_steps) + except (TypeError, ValueError): + warnings.warn( + "dtensor_cfg.clear_cache_every_n_steps must be an int; disabling CUDA cache clearing.", + RuntimeWarning, + ) + empty_cache_steps = None + if empty_cache_steps is not None and empty_cache_steps <= 0: + empty_cache_steps = None + if empty_cache_steps and not getattr(self, "_warned_cache_clear_overhead", False): + warnings.warn( + f"Emptying cache every {empty_cache_steps} microbatches; unnecessary use can incur large performance overhead." + ) + self._warned_cache_clear_overhead = True
ð§đ Nitpick comments (4)
nemo_rl/models/policy/dtensor_policy_worker_v2.py (2)
589-591: Avoid clearing at mb_idx 0; clear after each N microbatches.Current modulo clears on the first microbatch. Shift to (mb_idx + 1) % N == 0 so the first clear happens after N microbatches.
- if empty_cache_steps and mb_idx % empty_cache_steps == 0: + if empty_cache_steps and (mb_idx + 1) % empty_cache_steps == 0: torch.cuda.empty_cache()
581-584: Fix minor typos and improve wording.
- "unnnecessarily" â "unnecessarily"
- "Conditioanlly" â "Conditionally"
- Slightly tighten the warning text.
- warnings.warn( - f"Emptying cache every {empty_cache_steps} microbatches, doing so unnnecessarily would incur a large performance overhead." - ) + warnings.warn( + f"Emptying cache every {empty_cache_steps} microbatches; unnecessary use can incur large performance overhead." + ) @@ - # Conditioanlly empty cache when sensitive to fragmentation + # Conditionally empty cache when sensitive to fragmentationAlso applies to: 588-590
nemo_rl/models/policy/dtensor_policy_worker.py (2)
645-647: Shift clearing to occur after N microbatches, not at the first.- if empty_cache_steps and mb_idx % empty_cache_steps == 0: + if empty_cache_steps and (mb_idx + 1) % empty_cache_steps == 0: torch.cuda.empty_cache()
637-639: Typos in user-facing text and comment.
- "unnnecessarily" â "unnecessarily"
- "Conditioanlly" â "Conditionally"
- warnings.warn( - f"Emptying cache every {empty_cache_steps} microbatches, doing so unnnecessarily would incur a large performance overhead." - ) + warnings.warn( + f"Emptying cache every {empty_cache_steps} microbatches; unnecessary use can incur large performance overhead." + ) @@ - # Conditioanlly empty cache when sensitive to fragmentation + # Conditionally empty cache when sensitive to fragmentationAlso applies to: 644-645
ð Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
ð Files selected for processing (2)
nemo_rl/models/policy/dtensor_policy_worker.py(1 hunks)nemo_rl/models/policy/dtensor_policy_worker_v2.py(1 hunks)
ð Additional comments (1)
nemo_rl/models/policy/dtensor_policy_worker.py (1)
633-636: Consistent use ofclear_cache_every_n_stepsacross code, schemas, and configs. No occurrences ofempty_cache_every_n_stepswere found in the schema (nemo_rl/models/policy/__init__.py), policy workers (dtensor_policy_worker.py&_v2.py), example YAMLs, or docs.
Signed-off-by: Brian Yu <bxyu@nvidia.com>
|
retrying, failure is related to a race condition in mbridge. fyi @yfw https://github.com/NVIDIA-NeMo/RL/actions/runs/17603234870/job/50009498530#step:3:8329 |
Signed-off-by: Brian Yu <bxyu@nvidia.com>
Signed-off-by: Brian Yu <bxyu@nvidia.com>
3 participants
What does this PR do ?
There is a mismatch between the variable name to
clear_cache_every_n_stepsand what is actually used in the code which isempty_cache_every_n_steps. I noticed this when some jobs that weren't oom'ing last week started oom'ing this week and it's because we stopped clearing the cache here.Issues
List issues that this PR closes (syntax):
Usage
# Add a code snippet demonstrating how to use thisBefore your PR is "Ready for review"
Pre checks:
Additional Information
Summary by CodeRabbit