feat: add DP memory controls and physical batch sizing

[binaryaaron]

Summary

Adds OOM-mitigation controls and observability to the DP training path for the transformers v5 upgrade.

• Memory controls config: new training.memory (TrainingMemoryControls) replaces the old NSS_* env-var knobs with first-class, validated fields -- disable_dp_bf16, chunked_causal_lm_loss, chunked_causal_lm_loss_tokens, debug_loss_memory. Plumbed into OpacusDPTrainer.
• Physical batch sizing: training.max_physical_batch_size caps the per-device microbatch and gradient_accumulation_steps="auto" derives accumulation from the logical batch. resolve_batching() returns a single ResolvedTrainingBatching the backend consumes; DP guarantees hold across accumulation steps.
• DP grad-sample mode: new differential_privacy.grad_sample_mode (hooks | ghost) selecting Opacus fast vs ghost gradient clipping; opacus pinned to ==1.6.0.
• Preflight VRAM: the headroom check estimates from the resolved quantization state (QLoRA without quantize_model uses unquantized VRAM) instead of assuming quantized weights.
• Training observability: a schema-frozen TrainingObservability event (peak VRAM via NvmlPeakSampler, peak loss-logits, resolved batching, grad-sample mode) is emitted once per fit to structured logs and mirrored to wandb through a generic log_observability_event sink (WandbLoggable Protocol). Shared NVML/loadavg sampling primitives now live in observability.py. The opt-in causal-LM loss probe installs/tears down symmetrically in train(); the trainer stashes the event for the backend to forward, keeping the privacy layer independent of the CLI.

Coordination

The generic log_observability_event sink and the promoted NVML primitives are intended to converge with the generation-side observability work (#526). #526 lands first as the canonical home; this branch rebases and drops duplicates afterward.

Test plan

• mise run check (ruff + format + copyright + ty)
• Unit: config validation (TrainingMemoryControls, batching), configurator CLI flags, preflight VRAM, observability event/sink
• DP CPU smoke: one-step (hooks + ghost), gradient-accumulation (hooks + ghost), loss-probe install/uninstall round-trip
• GPU integration / e2e (run before merge)

Summary by CodeRabbit

Release Notes

• New Features

  • Added auto-mode support for gradient accumulation with resolved batching, including optional physical batch size capping.
  • Introduced differential privacy per-sample gradient modes: "hooks" and "ghost".
  • Added DP-focused training memory controls for chunked loss and loss-memory diagnostics.
  • Emit a structured training.complete analytics event including resolved batching and DP details.

• Documentation

  • Expanded configuration and out-of-memory troubleshooting, plus updated quantization guidance to training.quantization_scheme.

• Bug Fixes

  • Improved metric table formatting for per-second throughput values.

• Chores

  • Pinned Opacus to version 1.6.0.

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

Walkthrough

The PR introduces Opacus per-sample gradient mode selection ("hooks" vs "ghost"), physical batch size capping with automatic gradient accumulation derivation, training memory controls (DP-focused optimizations), and structured training-complete observability with peak VRAM/logits tracking. Opacus is pinned to 1.6.0 to support ghost clipping. Configuration, backend, trainer, and comprehensive test coverage are updated across all integration layers.

Changes

DP training memory control, batching resolution, and observability

Layer / File(s)	Summary
Training config types and batching resolution src/nemo_safe_synthesizer/config/training.py	Adds ResolvedTrainingBatching frozen dataclass with resolved per-device batch size, gradient accumulation steps, and effective batch size. Introduces TrainingMemoryControls for DP memory optimizations (disable_dp_bf16, chunked_causal_lm_loss with validated chunked_causal_lm_loss_tokens, debug_loss_memory). Extends TrainingHyperparams with max_physical_batch_size (auto-capable), memory field, updates gradient_accumulation_steps validation to allow "auto", adds resolve_batching() method computing physical microbatches under optional caps via largest-divisor selection while preserving effective batch size, and adds resolve_quantization_scheme() helper for legacy quantization_bits fallback.
Differential Privacy grad_sample_mode config src/nemo_safe_synthesizer/config/differential_privacy.py	Adds grad_sample_mode: Literal["hooks", "ghost"] field to DifferentialPrivacyHyperparams with default "hooks", documenting supported Opacus per-sample gradient computation modes.
Parameters validation approach update src/nemo_safe_synthesizer/config/parameters.py, src/nemo_safe_synthesizer/sdk/config_builder.py	Updates SafeSynthesizerParameters.from_params and _resolve_config to use model_validate(kwargs) instead of model_copy(update=...), ensuring auto and memory-control fields flow correctly through validation and override merging.
Training observability event schema src/nemo_safe_synthesizer/training/training_observability.py	Introduces TrainingObservability Pydantic model for strict training.complete event with optional measurements (peak VRAM, peak logits, resolved batching, grad sample mode) and to_wandb_payload(prefix) helper for namespaced WandB logging of non-None fields.
DP trainer with configurable grad sampling and loss probe src/nemo_safe_synthesizer/privacy/dp_transformers/dp_utils.py	Extends OpacusDPTrainer with grad_sample_mode and memory_controls parameters; implements ghost mode via version-gated dynamic loading of Opacus Fast/Ghost clipping classes; installs idempotent Transformers causal-LM loss memory probe to track peak fp32 logits; routes ghost mode through separate training-step helpers validating logits shape against logits_to_keep; samples peak VRAM via NvmlPeakSampler and emits structured observability event after training; broadens model unwrap logic to detect _module attribute for both grad-sample modes.
PRV accountant integration update src/nemo_safe_synthesizer/privacy/dp_transformers/privacy_args.py	Updates PRV accountant construction to use PoissonSubsampledGaussianMechanism and revised parameter wiring (delta_error=delta/1000, max_self_compositions); modifies epsilon computation calls in both standard and noise-multiplier search paths to include explicit target_delta and step list arguments.
HuggingFace backend DP and batching integration src/nemo_safe_synthesizer/training/huggingface_backend.py	Wires resolve_batching() into standard training args construction for both modes; passes grad_sample_mode and memory_controls to OpacusDPTrainer; delegates quantization scheme resolution to training.resolve_quantization_scheme(); applies DP-specific Trainer constraints (remove_unused_columns=False, max_grad_norm=0.0, no gradient checkpointing); conditionally disables Trainer bf16 when memory_controls.disable_dp_bf16 is set with warning; forwards trainer observability to WandB via log_observability_event after training in best-effort finally block.
Observability metric formatting and VRAM estimation refactoring src/nemo_safe_synthesizer/observability.py, src/nemo_safe_synthesizer/preflight/checks/environment.py	Updates _format_table_value to exclude per-second metric keys (*_per_sec, *_per_second) from percentage rendering. Refactors bytes_per_base_weight to early-return 2.0 when quantize_model disabled and compute bits-per-byte from resolved quantization scheme when enabled. Adjusts VRAMHeadroomCheck logging messages for parameter counting clarity.
Documentation: configuration, troubleshooting, and dependencies docs/user-guide/configuration.md, docs/user-guide/troubleshooting.md, docs/user-guide/running.md, pyproject.toml	Expands configuration reference with batching/memory fields and constraints. Updates OOM troubleshooting steps covering physical batch caps, ghost mode, chunked loss, loss memory probing, and max_vram_fraction tuning. Clarifies defaults and auto-resolution behavior for batching parameters. Updates quantization scheme examples from quantization_bits to quantization_scheme with legacy mapping documentation. Pins opacus==1.6.0 to both cpu and cu129 optional dependency groups.
Configuration and parameter resolution tests tests/config/test_parameters.py	Adds comprehensive tests for batching resolution under physical caps, memory control validation, grad_sample_mode acceptance/rejection via params and YAML, auto-typed field parsing, and helper assertions validating resolved batching and largest-divisor selection; tests preservation of effective batch size under capping.
CLI override tests for new config fields tests/cli/test_run.py	Tests that --training__gradient_accumulation_steps, --training__max_physical_batch_size, and --privacy__grad_sample_mode CLI flags flow through synthesis_overrides and resolve correctly into the final SafeSynthesizerParameters object.
Pydantic Click options tests for nested memory config tests/configurator/test_pydantic_click_options.py	Validates that nested training.memory sub-config is exposed as flattened CLI flags (training__memory__*) and that overrides correctly populate the nested dictionary in parsed results.
Training observability and WandB integration tests tests/training/test_training_observability.py	Tests TRAINING_COMPLETE_EVENT constant, TrainingObservability schema defaults/validation, to_wandb_payload() prefixing and None-filtering, log_observability_event() no-op when inactive and logging when active, and _emit_training_observability assembling events from trainer state with optional peak logits metrics.
Table rendering tests for metrics tests/observability/test_observability.py	Tests that per-second metric keys render as raw numeric strings rather than percentages in rich table output.
VRAM estimation with quantization toggle tests tests/preflight/test_preflight.py	Tests QLoRA VRAM headroom behavior when training.quantize_model toggles, verifying expected memory checks under both quantized and unquantized modes.
SDK builder privacy config validation tests tests/sdk/test_builder.py	Tests that SafeSynthesizer.with_differential_privacy() rejects invalid grad_sample_mode values via both dict-based and model config application paths.
DP trainer smoke tests with ghost clipping and probes tests/smoke/test_training_cpu.py	Adds DP trainer helpers (_dp_trainer, _privacy_args, _private_data_collator, _tiny_lora_model) centralizing test setup. Adds smoke tests for ghost clipping version gating (requires Opacus ≥1.6.0), ghost and hooks training with gradient accumulation validating Opacus substep callback invocation, loss probe install/uninstall round-trip with global restoration verification, and DP+LoRA checkpoint saving (PEFT adapter only, no wrapped model weights).
DP utils unit tests for loss probe and ghost clipping tests/training/test_dp_utils.py	Tests loss memory probe idempotency (rejects overlapping installs with cleanup guarantee), ghost clipping loss computation error handling when models ignore logits_to_keep constraint, _chunked_cross_entropy matching HuggingFace semantics with masked labels and normalization scaling, probe skipping in ghost mode with warning emission, and observability emission robustness when logger fails.
HuggingFace backend batching and DP tests tests/training/test_huggingface_backend.py	Tests that _build_base_training_args() forwards resolved batching and that DP trainer configuration correctly interprets physical batch caps and derives gradient accumulation; validates sampling probability computation using effective batch sizing and grad_sample_mode parameter forwarding; confirms observability event logging on training failure via log_observability_event.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~75 minutes

Possibly related PRs

• NVIDIA-NeMo/Safe-Synthesizer#528: Introduces physical batch size capping and auto gradient accumulation derivation that directly supports new DP run-config batching semantics.
• NVIDIA-NeMo/Safe-Synthesizer#483: Establishes training.quantization_scheme and resolve_quantization_scheme() foundation that this PR builds upon for both standard and DP training.
• NVIDIA-NeMo/Safe-Synthesizer#480: Concurrent preflight VRAM estimation refactoring; this PR adjusts bytes_per_base_weight and VRAMHeadroomCheck using new quantization scheme resolution.

Suggested reviewers

• kendrickb-nvidia
• mckornfield

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 60.32% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title 'feat: add DP memory controls and physical batch sizing' accurately summarizes the main changes: introduction of training memory controls and physical batch sizing configuration options for differential privacy.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch binaryaaron/transformers-v5-oom

---

_{Comment @coderabbitai help to get the list of available commands.}

[codecov]

Codecov Report

❌ Patch coverage is 80.27682% with 171 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...fe_synthesizer/privacy/dp_transformers/dp_utils.py	62.60%	89 Missing ⚠️
tests/smoke/test_training_cpu.py	22.22%	77 Missing ⚠️
...o_safe_synthesizer/training/huggingface_backend.py	83.33%	2 Missing ⚠️
src/nemo_safe_synthesizer/config/training.py	97.05%	1 Missing ⚠️
...ynthesizer/privacy/dp_transformers/privacy_args.py	66.66%	1 Missing ⚠️
tests/configurator/test_pydantic_click_options.py	93.75%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

[coderabbitai]

Actionable comments posted: 2

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: ccb2f659-e86a-4d5d-acd2-485856ffab2e

📥 Commits

Reviewing files that changed from the base of the PR and between 2a8f7a1 and 7967224.

⛔ Files ignored due to path filters (1)

• uv.lock is excluded by !**/*.lock, !uv.lock

📒 Files selected for processing (22)

• docs/user-guide/configuration.md
• docs/user-guide/troubleshooting.md
• pyproject.toml
• src/nemo_safe_synthesizer/cli/wandb_setup.py
• src/nemo_safe_synthesizer/config/differential_privacy.py
• src/nemo_safe_synthesizer/config/parameters.py
• src/nemo_safe_synthesizer/config/training.py
• src/nemo_safe_synthesizer/observability.py
• src/nemo_safe_synthesizer/preflight/checks/environment.py
• src/nemo_safe_synthesizer/privacy/dp_transformers/dp_utils.py
• src/nemo_safe_synthesizer/sdk/config_builder.py
• src/nemo_safe_synthesizer/training/huggingface_backend.py
• src/nemo_safe_synthesizer/training/training_observability.py
• tests/cli/test_run.py
• tests/config/test_parameters.py
• tests/configurator/test_pydantic_click_options.py
• tests/observability/test_observability.py
• tests/preflight/test_preflight.py
• tests/sdk/test_builder.py
• tests/smoke/test_training_cpu.py
• tests/training/test_huggingface_backend.py
• tests/training/test_training_observability.py

📜 Review details

🧰 Additional context used

📓 Path-based instructions (19)

**/*.{md,markdown,py}

📄 CodeRabbit inference engine (.cursor/rules/agent-markdown-style.mdc)

**/*.{md,markdown,py}: Avoid decorative bold (**text**) in list items, body text, and docstrings; use structural cues (headers, list markers, colons, backticks) for emphasis instead
Use backticks for code identifiers, paths, and CLI commands in markdown and docstrings

Files:

• src/nemo_safe_synthesizer/training/training_observability.py
• tests/configurator/test_pydantic_click_options.py
• tests/sdk/test_builder.py
• src/nemo_safe_synthesizer/sdk/config_builder.py
• src/nemo_safe_synthesizer/cli/wandb_setup.py
• src/nemo_safe_synthesizer/training/huggingface_backend.py
• tests/training/test_huggingface_backend.py
• docs/user-guide/troubleshooting.md
• tests/preflight/test_preflight.py
• src/nemo_safe_synthesizer/config/differential_privacy.py
• tests/observability/test_observability.py
• tests/training/test_training_observability.py
• src/nemo_safe_synthesizer/config/parameters.py
• src/nemo_safe_synthesizer/observability.py
• docs/user-guide/configuration.md
• src/nemo_safe_synthesizer/preflight/checks/environment.py
• tests/smoke/test_training_cpu.py
• src/nemo_safe_synthesizer/config/training.py
• tests/cli/test_run.py
• src/nemo_safe_synthesizer/privacy/dp_transformers/dp_utils.py
• tests/config/test_parameters.py

**/*.py

📄 CodeRabbit inference engine (AGENTS.md)

**/*.py: Place durable implementation guidance in function and class docstrings for public contracts and source comments for local invariants
Target Python 3.11–3.13 with modern syntax (X | Y, list[str], Self). Python 3.14+ is not supported

**/*.py: Use NSSBaseModel for config/parameter models in config/ which define the user-facing configuration of NSS. Use raw BaseModel or module-specific bases for data transfer objects and internal structures.
Use BaseSettings for env/CLI settings. Prefer AliasChoices on individual fields for both Python name and env var name mapping.
Always include Field(description=...) for Pydantic model fields as the canonical field docstring
Use assignment-style Field() for Pydantic model fields (e.g., field_name = Field(default=..., description=...)) instead of Annotated when default, default_factory, or alias are used
Use Annotated only when the field carries additional metadata beyond Field() such as validators, AutoParam, DependsOnValidator, or constrained type aliases
For mutable defaults in dataclasses, use field(default_factory=list) never = []
Use StrEnum for string-valued enums used in configs/serialization. Use plain Enum for internal-only named constants.
Use observability.get_logger(__name__) for logging, never logging.getLogger() or structlog.get_logger() directly
Use category loggers: .runtime for internals, .user for progress/results, .system for system events
Never use print() for operational output. Use click.echo() for CLI output or sys.stdout.write() for raw output in tools.
Use extra={} in logger calls for data that downstream tools should query or aggregate (metrics, counts, durations). Use f-strings for human-readable context that doesn't need machine parsing.
Raise errors from the custom hierarchy with dual inheritance: SafeSynthesizerError, UserError, DataError, ParameterError, GenerationError, InternalError
Target Python 3.11–3.13; do not us...

Files:

• src/nemo_safe_synthesizer/training/training_observability.py
• tests/configurator/test_pydantic_click_options.py
• tests/sdk/test_builder.py
• src/nemo_safe_synthesizer/sdk/config_builder.py
• src/nemo_safe_synthesizer/cli/wandb_setup.py
• src/nemo_safe_synthesizer/training/huggingface_backend.py
• tests/training/test_huggingface_backend.py
• tests/preflight/test_preflight.py
• src/nemo_safe_synthesizer/config/differential_privacy.py
• tests/observability/test_observability.py
• tests/training/test_training_observability.py
• src/nemo_safe_synthesizer/config/parameters.py
• src/nemo_safe_synthesizer/observability.py
• src/nemo_safe_synthesizer/preflight/checks/environment.py
• tests/smoke/test_training_cpu.py
• src/nemo_safe_synthesizer/config/training.py
• tests/cli/test_run.py
• src/nemo_safe_synthesizer/privacy/dp_transformers/dp_utils.py
• tests/config/test_parameters.py

src/nemo_safe_synthesizer/**/*.py

📄 CodeRabbit inference engine (CONTRIBUTING.md)

src/nemo_safe_synthesizer/**/*.py: Source code must remain Python 3.11 syntax-compatible. Do not use Python 3.12-only syntax such as PEP 695 type statements or bracketed generic class/function parameters in shared package code
Write Python docstrings in Google style format so they auto-generate into API reference pages via mkdocstrings and gen-files plugins

Files:

• src/nemo_safe_synthesizer/training/training_observability.py
• src/nemo_safe_synthesizer/sdk/config_builder.py
• src/nemo_safe_synthesizer/cli/wandb_setup.py
• src/nemo_safe_synthesizer/training/huggingface_backend.py
• src/nemo_safe_synthesizer/config/differential_privacy.py
• src/nemo_safe_synthesizer/config/parameters.py
• src/nemo_safe_synthesizer/observability.py
• src/nemo_safe_synthesizer/preflight/checks/environment.py
• src/nemo_safe_synthesizer/config/training.py
• src/nemo_safe_synthesizer/privacy/dp_transformers/dp_utils.py

**/*.{py,sh,yaml,yml,md}

📄 CodeRabbit inference engine (CONTRIBUTING.md)

All source files (.py, .sh, .yaml, .yml, .md) must include SPDX copyright headers. Use mise run format to add them automatically

Files:

• src/nemo_safe_synthesizer/training/training_observability.py
• tests/configurator/test_pydantic_click_options.py
• tests/sdk/test_builder.py
• src/nemo_safe_synthesizer/sdk/config_builder.py
• src/nemo_safe_synthesizer/cli/wandb_setup.py
• src/nemo_safe_synthesizer/training/huggingface_backend.py
• tests/training/test_huggingface_backend.py
• docs/user-guide/troubleshooting.md
• tests/preflight/test_preflight.py
• src/nemo_safe_synthesizer/config/differential_privacy.py
• tests/observability/test_observability.py
• tests/training/test_training_observability.py
• src/nemo_safe_synthesizer/config/parameters.py
• src/nemo_safe_synthesizer/observability.py
• docs/user-guide/configuration.md
• src/nemo_safe_synthesizer/preflight/checks/environment.py
• tests/smoke/test_training_cpu.py
• src/nemo_safe_synthesizer/config/training.py
• tests/cli/test_run.py
• src/nemo_safe_synthesizer/privacy/dp_transformers/dp_utils.py
• tests/config/test_parameters.py

**/*

⚙️ CodeRabbit configuration file

**/*: Review as a senior maintainer for NeMo Safe Synthesizer. Prioritize issues that can change behavior, break user workflows, weaken privacy guarantees, hide failures, make tests unreliable, or create maintenance risk. Avoid generic style commentary unless it points to a concrete project convention that automated tools will not catch.
Comment only when the finding is actionable and tied to changed code. For each finding, state the impact, the condition that triggers it, and the smallest practical fix. Prefer one precise comment over broad advice. Do not ask for refactors outside the PR scope unless the changed code creates the problem.
Review type guidance: - Potential issue: use for correctness bugs, data loss, privacy leaks,
security risks, broken public APIs, invalid config behavior, missing
validation, hidden failures, nondeterministic tests, or CI breakage.

• Refactor suggestion: use for local maintainability problems introduced
  by the diff when they have clear future cost, such as duplicated setup,
  unclear boundaries, over-mocking, avoidable complexity, or opaque test
  helpers.
• Nitpick: avoid in chill mode. Do not emit formatting, import-order,
  wording, or style-only comments unless automated tools cannot catch the
  issue and it affects maintainability.

Severity guidance: - Critical: security/privacy leaks, data loss, training/test/holdout
contamination, or broken release/package/core pipeline execution.

• Major: incorrect generation/training/evaluation behavior, broken
  CLI/SDK public API, invalid config defaults or validators, or GPU/vLLM
  cleanup and process-isolation bugs likely to fail CI or production
  runs.
• Minor: localized bugs, missing focused tests for changed behavior, or
  bad test patterns that weaken regression coverage.
• Trivial: small cleanup with no behavior impact. Usually suppress in
  chill mode.
• Info: context only. Avoid unless it helps reviewers understand risk.
  Safe-Synthesizer-specific review focus: - Data ...

Files:

• src/nemo_safe_synthesizer/training/training_observability.py
• tests/configurator/test_pydantic_click_options.py
• tests/sdk/test_builder.py
• src/nemo_safe_synthesizer/sdk/config_builder.py
• src/nemo_safe_synthesizer/cli/wandb_setup.py
• src/nemo_safe_synthesizer/training/huggingface_backend.py
• tests/training/test_huggingface_backend.py
• docs/user-guide/troubleshooting.md
• tests/preflight/test_preflight.py
• pyproject.toml
• src/nemo_safe_synthesizer/config/differential_privacy.py
• tests/observability/test_observability.py
• tests/training/test_training_observability.py
• src/nemo_safe_synthesizer/config/parameters.py
• src/nemo_safe_synthesizer/observability.py
• docs/user-guide/configuration.md
• src/nemo_safe_synthesizer/preflight/checks/environment.py
• tests/smoke/test_training_cpu.py
• src/nemo_safe_synthesizer/config/training.py
• tests/cli/test_run.py
• src/nemo_safe_synthesizer/privacy/dp_transformers/dp_utils.py
• tests/config/test_parameters.py

src/nemo_safe_synthesizer/training/**/*.py

⚙️ CodeRabbit configuration file

Review training changes for dataset preprocessing, model path handling, artifact writes, LoRA/DP behavior, GPU memory usage, reproducibility, and cleanup on failure.

Files:

• src/nemo_safe_synthesizer/training/training_observability.py
• src/nemo_safe_synthesizer/training/huggingface_backend.py

src/**/*.py

⚙️ CodeRabbit configuration file

Review library code against STYLE_GUIDE.md. Focus on behavior, API contracts, error handling, resource cleanup, typing, logging, and user-facing failures. Public APIs and nontrivial functions need Google-style docstrings.

Files:

• src/nemo_safe_synthesizer/training/training_observability.py
• src/nemo_safe_synthesizer/sdk/config_builder.py
• src/nemo_safe_synthesizer/cli/wandb_setup.py
• src/nemo_safe_synthesizer/training/huggingface_backend.py
• src/nemo_safe_synthesizer/config/differential_privacy.py
• src/nemo_safe_synthesizer/config/parameters.py
• src/nemo_safe_synthesizer/observability.py
• src/nemo_safe_synthesizer/preflight/checks/environment.py
• src/nemo_safe_synthesizer/config/training.py
• src/nemo_safe_synthesizer/privacy/dp_transformers/dp_utils.py

**

⚙️ CodeRabbit configuration file

**:

AGENTS.md

Guide for AI agents (Cursor, Windsurf, Claude Code, etc.) working in the Safe-Synthesizer repo.

This project loads local developer preferences from @AGENTS.local.md. You MUST read this file if it exists and give its instructions top priority.

Skills

Repo-specific skills live in .agents/skills/; see .agents/README.md for the catalog. Read a skill when the task matches its scope instead of copying workflow details into this file.

Durable implementation guidance belongs with the code it describes: function and class docstrings for public contracts and source comments for local invariants. Test-suite guidance belongs in tests/TESTING.md.

Repo Conventions

See STYLE_GUIDE.md for detailed code style conventions (Python, markdown, Dockerfiles, shell scripts, testing, config files, docstrings).

Use uv for everything -- never pip or raw python. Python 3.11–3.13 with modern syntax (X | Y, list[str], Self). Python 3.14+ is not supported.

Common commands: mise run test (unit tests), mise run format (auto-fix formatting + lint + copyright), mise run check (read-only local quality checks), mise run validate (pre-PR quality, lock, and CI unit checks), mise run typecheck (ty only). Always use mise tasks or the wrapper scripts in tools/ instead of running ruff or ty directly. Use uv run for Python execution. When in doubt, inspect mise tasks and pytest --markers.

The canonical uv sync command for a full GPU/dev environment is:

uv sync --frozen --extra cu129 --extra engine --group dev

Bare uv sync --frozen (without extras) installs an incomplete environment -- ty, import checks, and GPU tests will fail.

Feature branches off main. Branch names often include an issue number prefix (e.g., <author>/123-short-name).

Do ...

Files:

• src/nemo_safe_synthesizer/training/training_observability.py
• tests/configurator/test_pydantic_click_options.py
• tests/sdk/test_builder.py
• src/nemo_safe_synthesizer/sdk/config_builder.py
• src/nemo_safe_synthesizer/cli/wandb_setup.py
• src/nemo_safe_synthesizer/training/huggingface_backend.py
• tests/training/test_huggingface_backend.py
• docs/user-guide/troubleshooting.md
• tests/preflight/test_preflight.py
• pyproject.toml
• src/nemo_safe_synthesizer/config/differential_privacy.py
• tests/observability/test_observability.py
• tests/training/test_training_observability.py
• src/nemo_safe_synthesizer/config/parameters.py
• src/nemo_safe_synthesizer/observability.py
• docs/user-guide/configuration.md
• src/nemo_safe_synthesizer/preflight/checks/environment.py
• tests/smoke/test_training_cpu.py
• src/nemo_safe_synthesizer/config/training.py
• tests/cli/test_run.py
• src/nemo_safe_synthesizer/privacy/dp_transformers/dp_utils.py
• tests/config/test_parameters.py

**/test_*.py

📄 CodeRabbit inference engine (AGENTS.md)

Use the unit marker instead of the deprecated unit_test marker for test identification

Files:

• tests/configurator/test_pydantic_click_options.py
• tests/sdk/test_builder.py
• tests/training/test_huggingface_backend.py
• tests/preflight/test_preflight.py
• tests/observability/test_observability.py
• tests/training/test_training_observability.py
• tests/smoke/test_training_cpu.py
• tests/cli/test_run.py
• tests/config/test_parameters.py

tests/**/*.py

📄 CodeRabbit inference engine (tests/TESTING.md)

tests/**/*.py: Define pytest markers in pytest.ini with --strict-markers enabled. Use exactly one category marker (unit, smoke, e2e) per test. Category modifiers include slow for long-running tests, requires_gpu for CUDA-dependent tests, vllm for vLLM backend tests, smollm2 for SmolLM2 Hub download tests, and noautouse to skip autouse fixtures
For ParsedResponse mocking, use valid_records=[...], invalid_records=[...], errors=[...], and prompt_number=int. Use fixture_mock_processor or fixture_mock_processor_without_valid_records helpers
Use pytest.importorskip to gate tests on optional dependencies that require specific extras (e.g., sentence_transformers and vllm require cu129 extra)
For vLLM tests that call .generate(), mark with @pytest.mark.vllm, use per-file process isolation (-n 0), and create dedicated test:smoke:gpu:* mise tasks. vLLM pre-allocates all GPU memory and never releases it within a process, causing OOM in later tests if not isolated
Convert nullable columns to pd.Int64Dtype() or pd.BooleanDtype() before assigning np.nan to avoid pandas dtype conversion errors
When using Faker for test data generation, call fake.seed_instance(seed) and random.seed(seed) for reproducibility
Use print() statements in test functions for debug output (ruff T201 is suppressed for tests/). Import shared methods from conftest.py using relative imports when needed across test files (e.g., from .conftest import train_with_sdk). Do not import from other test files like tests/cli/helpers.py due to pytest discovery limitations

tests/**/*.py: File naming for tests: test_*.py; class naming: Test*; function naming: test_<module>_<expected_behavior>
Use fixture_ prefix convention for pytest fixtures for grep-ability; add a one-line docstring describing the fixture's purpose and data
Use function-scoped fixtures by default in pytest; session scope only when empirically justified by test runtime
Use...

Files:

• tests/configurator/test_pydantic_click_options.py
• tests/sdk/test_builder.py
• tests/training/test_huggingface_backend.py
• tests/preflight/test_preflight.py
• tests/observability/test_observability.py
• tests/training/test_training_observability.py
• tests/smoke/test_training_cpu.py
• tests/cli/test_run.py
• tests/config/test_parameters.py

---

⚙️ CodeRabbit configuration file

Review tests against tests/TESTING.md. Check marker usage, fixture naming, tmp_path usage, determinism, and GPU/vLLM process-isolation requirements. Flag slop tests that only check that code runs, assert result is not None when stronger invariants exist, over-mock internal implementation details, patch around the bug instead of reproducing it, or add broad snapshot/golden churn without a clear contract. Flag change detector tests that fail on harmless refactors, formatting, record ordering, incidental wording, or private implementation details without demonstrating a behavior regression. Prefer existing fixtures or focused new fixtures for repeated setup; keep tests DRY when reasonable without making the behavior under test opaque. print() is allowed in tests.

Files:

• tests/configurator/test_pydantic_click_options.py
• tests/sdk/test_builder.py
• tests/training/test_huggingface_backend.py
• tests/preflight/test_preflight.py
• tests/observability/test_observability.py
• tests/training/test_training_observability.py
• tests/smoke/test_training_cpu.py
• tests/cli/test_run.py
• tests/config/test_parameters.py

tests/**

📄 CodeRabbit inference engine (.cursor/rules/repo-navigation.mdc)

tests/**: Mirror src/ directory structure in tests/ directory for test organization
Auto-mark tests by directory: tests/e2e/ → e2e, tests/smoke/ → smoke, otherwise default to unit

Files:

• tests/configurator/test_pydantic_click_options.py
• tests/sdk/test_builder.py
• tests/training/test_huggingface_backend.py
• tests/preflight/test_preflight.py
• tests/observability/test_observability.py
• tests/training/test_training_observability.py
• tests/smoke/test_training_cpu.py
• tests/cli/test_run.py
• tests/config/test_parameters.py

---

⚙️ CodeRabbit configuration file

tests/**:

Testing Guide

Comprehensive testing reference for Safe-Synthesizer developers. Covers commands, markers, test data, fixtures, and gotchas.

Read First

1. tests/conftest.py -- auto-marking, load_test_dataset/load_test_dataframe, fixture_mock_processor pattern
2. pytest.ini -- markers, asyncio, timeout
3. tests/evaluation/conftest.py -- most complex: Faker-based make_df, nullable dtype conversion
4. tests/generation/conftest.py -- JSONL/schema fixtures, fixture_valid_iris_dataset_jsonl_and_schema

Running Tests

All mise test tasks, grouped by scope:

mise run test                              # Unit (excludes slow, e2e, and smoke)
mise run test:unit-slow                    # Unit tests including slow (excludes e2e and smoke)
mise run test:smoke                        # CPU smoke tests (~few min, no GPU required)
mise run test:smoke:gpu                    # All staged GPU smoke tests (requires CUDA)
mise run test:smoke:gpu:train-only
mise run test:smoke:gpu:generation
mise run test:smoke:gpu:resume
mise run test:smoke:gpu:structured-generation
mise run test:smoke:gpu:timeseries
mise run test:smoke:gpu:smollm2
mise run test:e2e                          # All e2e (requires CUDA) -- runs default + dp
mise run test:e2e:default                  # e2e default (no-DP) tests only
mise run test:e2e:dp                       # e2e DP tests only
mise run test:ci                           # CI unit tests with coverage (excludes slow, e2e, gpu, smoke)
mise run test:ci-slow                      # CI slow tests with coverage
mise run test:ci-container                 # CI tests in a Linux container (Docker/Podman)

Run a single test:

uv run --frozen pytest tests/path/test_file.py::test_name -vvs -n0

Test runner: uv run --frozen pytest -n auto --dist loadscope -vv...

Files:

• tests/configurator/test_pydantic_click_options.py
• tests/sdk/test_builder.py
• tests/training/test_huggingface_backend.py
• tests/preflight/test_preflight.py
• tests/observability/test_observability.py
• tests/training/test_training_observability.py
• tests/smoke/test_training_cpu.py
• tests/cli/test_run.py
• tests/config/test_parameters.py

**/*.{md,markdown}

📄 CodeRabbit inference engine (.cursor/rules/agent-markdown-style.mdc)

**/*.{md,markdown}: Bold is acceptable only in markdown tables where it's the conventional way to mark header-like cells in the body
Use ## headers to segment markdown sections instead of bold text
Use -- (em-dash) instead of - (hyphen) for asides in markdown

Files:

• docs/user-guide/troubleshooting.md
• docs/user-guide/configuration.md

docs/**/*.md

📄 CodeRabbit inference engine (.cursor/rules/writing-docs.mdc)

docs/**/*.md: Use MkDocs Material admonition syntax (!!! note, !!! warning, ??? tip) for highlighting important information and collapsible sections in documentation
Use MkDocs Material tabs syntax (=== "Label") to present alternative views or language-specific examples in documentation
Use code block syntax with title and highlight line parameters (title="filename", hl_lines="2 3") for code examples in documentation
Use Mermaid diagram syntax (```mermaid flowchart, etc.) for visualizations in documentation

Classify documentation pages in docs/ as tutorial, how-to, explanation, or reference per the Diataxis framework. Use MkDocs Material syntax: admonitions (!!! note), tabs (===), code blocks with titles and highlights.

Documentation markdown files must follow MkDocs Material Markdown extensions syntax including admonitions, content tabs, code blocks with annotations, and Mermaid diagrams

Files:

• docs/user-guide/troubleshooting.md
• docs/user-guide/configuration.md

**/*.md

📄 CodeRabbit inference engine (STYLE_GUIDE.md)

**/*.md: No decorative **bold** in markdown body text, list items, or docstrings. Use headers, list markers, colons, and backticks for structure.
Use -- (em-dash) for asides in markdown, not - (hyphen)
Use single backticks for code identifiers, paths, and CLI commands in markdown
For Mermaid diagrams: use no spaces in node IDs, quote labels with special characters, avoid explicit colors or styles

Markdown documentation must be checked for consistency and formatting. Run mise run check to validate markdown files

Files:

• docs/user-guide/troubleshooting.md
• docs/user-guide/configuration.md

docs/**

⚙️ CodeRabbit configuration file

Review documentation as MkDocs Material content. Check Diataxis fit, accurate commands, internal links, code fences, and markdown style from STYLE_GUIDE.md.

Files:

• docs/user-guide/troubleshooting.md
• docs/user-guide/configuration.md

pyproject.toml

📄 CodeRabbit inference engine (.cursor/rules/repo-navigation.mdc)

Configure package metadata, dependencies, extras (cpu/cu129/engine), and uv configuration in pyproject.toml

Order sections in pyproject.toml as: [project], [dependency-groups], [project.optional-dependencies], [tool.uv], [build-system], [tool.*]

Files:

• pyproject.toml

---

⚙️ CodeRabbit configuration file

Treat pyproject.toml as high-risk. Check package metadata, uv indexes, dependency groups, optional extras, Python version bounds, hatch config, ty config, script entry points, dependency consistency, and whether changes require regenerating uv.lock.

Files:

• pyproject.toml

**/*.toml

📄 CodeRabbit inference engine (STYLE_GUIDE.md)

**/*.toml: Use spaces around = for key-value pairs in TOML files
Use # comment in TOML files; use inline comments for dependency pins

Files:

• pyproject.toml

src/nemo_safe_synthesizer/config/**/*.py

⚙️ CodeRabbit configuration file

Treat config changes as user-facing API changes. Check Pydantic field descriptions, defaults, validators, aliases, override behavior, CLI help text impact, YAML compatibility, and documented parameter semantics.

Files:

• src/nemo_safe_synthesizer/config/differential_privacy.py
• src/nemo_safe_synthesizer/config/parameters.py
• src/nemo_safe_synthesizer/config/training.py

src/nemo_safe_synthesizer/privacy/**/*.py

⚙️ CodeRabbit configuration file

Treat privacy changes as high-risk. Check DP accounting, parameter validation, data leakage, seed handling, model state persistence, and whether privacy guarantees are documented accurately.

Files:

• src/nemo_safe_synthesizer/privacy/dp_transformers/dp_utils.py

🧠 Learnings (11)

📚 Learning: 2026-05-27T22:20:37.354Z

Learnt from: kendrickb-nvidia
Repo: NVIDIA-NeMo/Safe-Synthesizer PR: 520
File: tests/generation/test_vllm_backend.py:556-587
Timestamp: 2026-05-27T22:20:37.354Z
Learning: In NVIDIA-NeMo/Safe-Synthesizer, `tests/conftest.py`’s `pytest_collection_modifyitems` hook applies pytest category markers automatically based on each test file’s path: tests under `/e2e/` get `pytest.mark.e2e`, tests under `/smoke/` get `pytest.mark.smoke`, and all other tests get `pytest.mark.unit`. Therefore, when reviewing pytest tests outside `tests/e2e/` and `tests/smoke/`, do not flag missing explicit `pytest.mark.unit` decorators on test classes/functions as an issue (the hook will add them during collection). If a new test directory/category is introduced, ensure the hook is updated so it’s categorized correctly.

Applied to files:

• tests/configurator/test_pydantic_click_options.py
• tests/sdk/test_builder.py
• tests/training/test_huggingface_backend.py
• tests/preflight/test_preflight.py
• tests/observability/test_observability.py
• tests/training/test_training_observability.py
• tests/smoke/test_training_cpu.py
• tests/cli/test_run.py
• tests/config/test_parameters.py

📚 Learning: 2026-06-03T23:08:50.142Z

Learnt from: CR
Repo: NVIDIA-NeMo/Safe-Synthesizer PR: 0
File: tests/TESTING.md:0-0
Timestamp: 2026-06-03T23:08:50.142Z
Learning: Applies to tests/**/*.py : For `ParsedResponse` mocking, use `valid_records=[...]`, `invalid_records=[...]`, `errors=[...]`, and `prompt_number=int`. Use `fixture_mock_processor` or `fixture_mock_processor_without_valid_records` helpers

Applied to files:

• tests/sdk/test_builder.py

📚 Learning: 2026-06-05T23:48:32.983Z

Learnt from: CR
Repo: NVIDIA-NeMo/Safe-Synthesizer PR: 0
File: STYLE_GUIDE.md:0-0
Timestamp: 2026-06-05T23:48:32.983Z
Learning: Applies to **/*.py : Use `NSSBaseModel` for config/parameter models in `config/` which define the user-facing configuration of NSS. Use raw `BaseModel` or module-specific bases for data transfer objects and internal structures.

Applied to files:

• src/nemo_safe_synthesizer/sdk/config_builder.py

📚 Learning: 2026-06-05T23:48:32.983Z

Learnt from: CR
Repo: NVIDIA-NeMo/Safe-Synthesizer PR: 0
File: STYLE_GUIDE.md:0-0
Timestamp: 2026-06-05T23:48:32.983Z
Learning: Applies to **/*.py : Use `Protocol` for structural subtyping when you need duck-typing boundaries in Python

Applied to files:

• src/nemo_safe_synthesizer/cli/wandb_setup.py

📚 Learning: 2026-06-03T23:09:02.641Z

Learnt from: CR
Repo: NVIDIA-NeMo/Safe-Synthesizer PR: 0
File: .cursor/rules/repo-navigation.mdc:0-0
Timestamp: 2026-06-03T23:09:02.641Z
Learning: Applies to pyproject.toml : Configure package metadata, dependencies, extras (cpu/cu129/engine), and uv configuration in `pyproject.toml`

Applied to files:

• pyproject.toml

📚 Learning: 2026-06-05T23:48:32.983Z

Learnt from: CR
Repo: NVIDIA-NeMo/Safe-Synthesizer PR: 0
File: STYLE_GUIDE.md:0-0
Timestamp: 2026-06-05T23:48:32.983Z
Learning: Applies to **/*.py : Use `Annotated` only when the field carries additional metadata beyond `Field()` such as validators, `AutoParam`, `DependsOnValidator`, or constrained type aliases

Applied to files:

• src/nemo_safe_synthesizer/config/differential_privacy.py

📚 Learning: 2026-06-03T23:08:50.142Z

Learnt from: CR
Repo: NVIDIA-NeMo/Safe-Synthesizer PR: 0
File: tests/TESTING.md:0-0
Timestamp: 2026-06-03T23:08:50.142Z
Learning: Applies to tests/**/*.py : Use `print()` statements in test functions for debug output (ruff `T201` is suppressed for `tests/`). Import shared methods from `conftest.py` using relative imports when needed across test files (e.g., `from .conftest import train_with_sdk`). Do not import from other test files like `tests/cli/helpers.py` due to pytest discovery limitations

Applied to files:

• tests/smoke/test_training_cpu.py

📚 Learning: 2026-06-03T23:08:50.142Z

Learnt from: CR
Repo: NVIDIA-NeMo/Safe-Synthesizer PR: 0
File: tests/TESTING.md:0-0
Timestamp: 2026-06-03T23:08:50.142Z
Learning: Applies to tests/**/conftest.py : Define test data fixtures with the `fixture_` prefix (e.g., `fixture_iris_dataset`, `fixture_dow_jones_index_dataset`, `fixture_clinc_oos_dataset`). Store stub datasets in `tests/stub_datasets/` and use helper functions `load_test_dataset(filename)` and `load_test_dataframe(filename)` from root `conftest.py` to load test data

Applied to files:

• tests/smoke/test_training_cpu.py

📚 Learning: 2026-06-04T16:14:09.868Z

Learnt from: binaryaaron
Repo: NVIDIA-NeMo/Safe-Synthesizer PR: 526
File: tests/generation/test_vllm_backend.py:399-509
Timestamp: 2026-06-04T16:14:09.868Z
Learning: In this repo, only apply `pytest.mark.vllm` to smoke tests under `tests/smoke/` that actually run real vLLM GPU generation and therefore require per-file process isolation (e.g., `test-smoke-gpu-*` Makefile targets). Do not apply `pytest.mark.vllm` to unit-style tests under `tests/generation/` that merely import `vllm_backend` but never instantiate a real vLLM engine and never call `.generate()` (GPU not required). Note that `tests/conftest.py` auto-marks these as `unit` via `pytest_collection_modifyitems`, and `vllm` is not among the auto-mark categories—so if a test in `tests/generation/` has `vllm`, it should be treated as a review issue unless it meets the real GPU generation criteria above.

Applied to files:

• tests/smoke/test_training_cpu.py

📚 Learning: 2026-06-05T23:48:32.983Z

Learnt from: CR
Repo: NVIDIA-NeMo/Safe-Synthesizer PR: 0
File: STYLE_GUIDE.md:0-0
Timestamp: 2026-06-05T23:48:32.983Z
Learning: Applies to **/*.py : Add `from __future__ import annotations` to every Python module to make all annotations strings consistently

Applied to files:

• src/nemo_safe_synthesizer/config/training.py
• src/nemo_safe_synthesizer/privacy/dp_transformers/dp_utils.py

📚 Learning: 2026-06-01T18:13:50.774Z

Learnt from: binaryaaron
Repo: NVIDIA-NeMo/Safe-Synthesizer PR: 538
File: src/nemo_safe_synthesizer/cli/utils.py:0-0
Timestamp: 2026-06-01T18:13:50.774Z
Learning: In `src/nemo_safe_synthesizer/cli/utils.py`, `_propagate_runtime_settings_to_env` intentionally uses write-only (non-`None`) propagation. The function runs exactly once per `common_setup` call in a one-shot CLI process, so there is no prior-invocation leakage risk. A clear-on-`None` guard (else: os.environ.pop(...)) would also be ineffective because `CLISettings` re-reads `os.environ` at construction time, meaning any value written by a hypothetical prior run is already reflected as non-`None` in the field — the else branch would never fire. A real cross-run isolation fix would require snapshotting/restoring `os.environ` around each run, which is out of scope for this PR.

Applied to files:

• src/nemo_safe_synthesizer/privacy/dp_transformers/dp_utils.py

🪛 Ruff (0.15.15)

tests/smoke/test_training_cpu.py

[warning] 182-182: Pattern passed to match= contains metacharacters but is neither escaped nor raw

(RUF043)

src/nemo_safe_synthesizer/privacy/dp_transformers/dp_utils.py

[warning] 215-215: Do not call setattr with a constant attribute value. It is not any safer than normal property access.

Replace setattr with assignment

(B010)

---

[warning] 241-241: Do not call setattr with a constant attribute value. It is not any safer than normal property access.

Replace setattr with assignment

(B010)

---

[error] 891-892: try-except-pass detected, consider logging the exception

(S110)

🔇 Additional comments (18)

src/nemo_safe_synthesizer/preflight/checks/environment.py (1)

194-229: LGTM!

Also applies to: 392-400

tests/preflight/test_preflight.py (1)

260-282: LGTM!

docs/user-guide/configuration.md (1)

71-77: LGTM!

Also applies to: 222-223

docs/user-guide/troubleshooting.md (1)

157-179: LGTM!

Also applies to: 439-465

src/nemo_safe_synthesizer/config/training.py (1)

7-7: LGTM!

Also applies to: 35-40, 157-171, 174-230, 266-303, 501-538

src/nemo_safe_synthesizer/config/differential_privacy.py (1)

8-9: LGTM!

Also applies to: 78-90

src/nemo_safe_synthesizer/config/parameters.py (1)

217-221: LGTM!

src/nemo_safe_synthesizer/sdk/config_builder.py (1)

9-9: LGTM!

Also applies to: 122-126

tests/config/test_parameters.py (1)

12-37: LGTM!

Also applies to: 46-47, 50-211

tests/cli/test_run.py (1)

844-847: LGTM!

Also applies to: 902-905, 937-957

tests/configurator/test_pydantic_click_options.py (1)

483-512: LGTM!

tests/sdk/test_builder.py (1)

9-12: LGTM!

Also applies to: 197-208

src/nemo_safe_synthesizer/cli/wandb_setup.py (2)

278-278: Already flagged: remove the Ellipsis default from the Protocol method parameter.

---

14-15: LGTM!

Also applies to: 268-277, 281-302

src/nemo_safe_synthesizer/training/training_observability.py (1)

1-84: LGTM!

src/nemo_safe_synthesizer/observability.py (1)

84-85: LGTM!

Also applies to: 212-216, 1016-1149

tests/training/test_training_observability.py (1)

1-135: LGTM!

tests/observability/test_observability.py (1)

298-314: LGTM!

[greptile-apps]

Greptile Summary

This PR adds OOM-mitigation controls and structured observability to the DP training path for the transformers v5 upgrade. It replaces ad-hoc NSS_* env vars with first-class TrainingMemoryControls, introduces physical-batch sizing via resolve_batching(), adds a ghost-gradient clipping mode, fixes the VRAM preflight estimate for QLoRA-without-quantize_model, and emits a schema-frozen training.complete observability event per fit.

• Memory controls & batching: TrainingMemoryControls (disable_dp_bf16, chunked_causal_lm_loss, debug_loss_memory) flows into OpacusDPTrainer; resolve_batching() uses _largest_divisor_at_most to find a physical microbatch that divides the configured effective batch size exactly, preserving DP lot-size accounting across accumulation steps.
• Ghost clipping: grad_sample_mode="ghost" loads GradSampleModuleFastGradientClipping / DPLossFastGradientClipping from Opacus 1.6.0 via guarded dynamic import; a custom no_sync context manager bridges the Trainer API.
• Observability: TrainingObservability (peak VRAM, peak logits GiB bucket, resolved batching, grad-sample mode) is emitted from OpacusDPTrainer.train() and mirrored to wandb via the new log_observability_event / WandbLoggable sink in wandb_setup.py.

Confidence Score: 5/5

Safe to merge; training correctness and DP accounting are preserved across both gradient modes.

The DP accounting invariant (effective batch = physical batch × accumulation steps) is correctly preserved by resolve_batching(), and both the hooks and ghost clipping paths are covered by CPU smoke tests. The two findings are observability-only: the peak_loss_logits_gb_bucket_le metric over-reports when chunking is active, and the opacus version spec is looser than stated in the PR description. Neither affects training correctness or DP guarantees.

src/nemo_safe_synthesizer/privacy/dp_transformers/dp_utils.py — the _record_peak_loss_logits call and opacus version spec in pyproject.toml are worth a second look before the next uv lock upgrade.

Important Files Changed

Filename	Overview
src/nemo_safe_synthesizer/config/training.py	Adds TrainingMemoryControls, ResolvedTrainingBatching, and resolve_batching(); the divisor-finding logic correctly preserves effective batch size for DP accounting.
src/nemo_safe_synthesizer/privacy/dp_transformers/dp_utils.py	Large addition of ghost clipping path, chunked cross-entropy, and training observability; peak_loss_logits recording over-reports when chunking is active (records full fp32 size regardless).
src/nemo_safe_synthesizer/privacy/dp_transformers/privacy_args.py	Migrates from old Accountant API to new PRVAccountant/PoissonSubsampledGaussianMechanism API; compute_epsilon call sites now pass delta and compositions list.
src/nemo_safe_synthesizer/training/training_observability.py	New schema-frozen TrainingObservability event with extra="forbid"; clean design with Protocol-based WandbLoggable decoupling.
src/nemo_safe_synthesizer/training/huggingface_backend.py	Wires resolve_batching() into TrainingArguments and threads grad_sample_mode + memory_controls into OpacusDPTrainer; observability forwarded in finally block.
src/nemo_safe_synthesizer/preflight/checks/environment.py	VRAM estimate now uses resolve_batching() for physical batch size and resolve_quantization_scheme() for QLoRA-without-quantize_model accuracy fix.
pyproject.toml	Opacus bumped to >=1.6.0 (resolves to 1.6.0 in lockfile); ghost clipping imports from internal Opacus module paths that may change in future versions.
src/nemo_safe_synthesizer/config/differential_privacy.py	Adds grad_sample_mode field with Literal["hooks", "ghost"] default "hooks"; clean addition.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant HFB as HuggingFaceBackend
    participant TB as TrainingHyperparams
    participant DPT as OpacusDPTrainer
    participant Ghost as GradSampleModuleFastGradientClipping
    participant Hooks as GradSampleModule (hooks)
    participant Obs as TrainingObservability
    participant WB as wandb

    HFB->>TB: resolve_batching()
    TB-->>HFB: ResolvedTrainingBatching
    HFB->>DPT: __init__(grad_sample_mode, memory_controls)
    alt "grad_sample_mode == hooks"
        DPT->>Hooks: wrap model
    else "grad_sample_mode == ghost"
        DPT->>Ghost: wrap model (via _load_ghost_clipping_classes)
    end
    HFB->>DPT: train()
    DPT->>DPT: NvmlPeakSampler (context)
    loop training steps
        alt hooks mode
            DPT->>DPT: _compute_hooks_causal_lm_loss()
        else ghost mode
            DPT->>DPT: _compute_ghost_clipping_loss()
        end
    end
    DPT->>Obs: _emit_training_observability(peak_vram_gb)
    DPT-->>HFB: last_training_observability
    HFB->>WB: "log_observability_event(event, prefix=training)"

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant HFB as HuggingFaceBackend
    participant TB as TrainingHyperparams
    participant DPT as OpacusDPTrainer
    participant Ghost as GradSampleModuleFastGradientClipping
    participant Hooks as GradSampleModule (hooks)
    participant Obs as TrainingObservability
    participant WB as wandb

    HFB->>TB: resolve_batching()
    TB-->>HFB: ResolvedTrainingBatching
    HFB->>DPT: __init__(grad_sample_mode, memory_controls)
    alt "grad_sample_mode == hooks"
        DPT->>Hooks: wrap model
    else "grad_sample_mode == ghost"
        DPT->>Ghost: wrap model (via _load_ghost_clipping_classes)
    end
    HFB->>DPT: train()
    DPT->>DPT: NvmlPeakSampler (context)
    loop training steps
        alt hooks mode
            DPT->>DPT: _compute_hooks_causal_lm_loss()
        else ghost mode
            DPT->>DPT: _compute_ghost_clipping_loss()
        end
    end
    DPT->>Obs: _emit_training_observability(peak_vram_gb)
    DPT-->>HFB: last_training_observability
    HFB->>WB: "log_observability_event(event, prefix=training)"

Loading

_{Reviews (13): Last reviewed commit: "test(training): cover ghost loss ignore ..." | Re-trigger Greptile}

[binaryaaron]

Let's land the vllm observability pr first.

[mckornfield]

just bumped some of andre's comments, but lg for my old ones

[mckornfield]

curious, did we have an effort to do structured logging/ is that applicable?

[binaryaaron]

agent (review-pr): Structured logging is applicable here

get_logger(__name__) already routes these messages through the structured logging stack, but the dictionaries are currently interpolated into the event text rather than emitted as structured context. Using extra={"ctx": ...} would make the fields queryable. This is a valid cleanup; it does not affect the configuration behavior changed here.

[mckornfield]

super nit, should we lift these guys up, since this is also in the type check? or is it here because it's expensive?

[binaryaaron]

agent (review-pr): The local import avoids a module cycle

This placement is intentional. configurator.parameters imports ParameterSchema from parameter_paths at module scope, so importing Parameters back at the top of parameter_paths would create a circular import. TYPE_CHECKING supplies the annotation, and the local runtime import defers the reverse dependency until the helper is called.

[mckornfield]

can we address this guy?

[mckornfield]

bump here too

[greptile-apps]

Ignored tokens dilute loss

Ghost mode keeps a token column when any sequence in the batch has a valid label there. In a mixed-length batch, a shorter sequence can still have -100 at one of those kept columns. That ignored position is passed into the Opacus loss reduction as a zero loss, so it is counted in the per-sequence mean instead of being excluded. The hooks path divides each sequence by its own valid-token count, but this path can shrink the loss and gradients for shorter examples, changing the clipping behavior for the same DP unit. This is reachable with batches where examples have different label masks.

Context Used: AGENTS.md (source)

[mckornfield]

handing out stamps