feat: add DP memory controls and physical batch sizing

docs/user-guide/configuration.md

@@ -142,8 +142,13 @@ for the full field list.
 | Field | Default | Description | Guidance |
 |-------|---------|-------------|----------|
 | `training.learning_rate` | `"auto"` | Initial learning rate for the `AdamW` optimizer. `"auto"` selects a model-specific default (Mistral: 1e-4, others: 5e-4) | Leave at `"auto"` for most cases; override with a float in (0, 1) to tune manually |
-| `training.batch_size` | `1` | Per-device batch size | Leave at 1; increase `gradient_accumulation_steps` for a larger effective batch |
-| `training.gradient_accumulation_steps` | `8` | Steps to accumulate before a backward pass; effective batch size = `batch_size` x this value | 8--32 typical |
+| `training.batch_size` | `1` | Per-device batch target. With integer `gradient_accumulation_steps`, effective batch size is `batch_size` x `gradient_accumulation_steps`; with `"auto"` accumulation, `batch_size` is the logical target | Leave at 1 for memory-constrained runs; increase only when VRAM headroom is available |
+| `training.max_physical_batch_size` | `"auto"` | Optional cap on the physical per-device microbatch sent to the Trainer. A numeric cap preserves the logical effective batch by increasing accumulation when possible | Use 1--4 to recover from activation OOMs while preserving the configured logical batch; leave `"auto"` otherwise |
+| `training.gradient_accumulation_steps` | `8` | Steps to accumulate before a backward pass; accepts an integer or `"auto"`. `"auto"` derives accumulation from `batch_size` and `max_physical_batch_size` | 8--32 typical for integer values; use `"auto"` with `max_physical_batch_size` for OOM triage |
+| `training.memory.disable_dp_bf16` | `false` | Disable the Trainer's bf16 autocast during DP training so model outputs stay in their original dtype. Only applies when DP is enabled | Enable to mitigate the fp32 logits-upcast memory spike at a speed cost |
+| `training.memory.chunked_causal_lm_loss` | `false` | Compute causal-LM cross entropy in token chunks instead of upcasting all logits to fp32 at once. DP path only | Enable for large-vocab OOMs at the loss upcast; small speed cost |
+| `training.memory.chunked_causal_lm_loss_tokens` | `1024` | Token chunk size used when `chunked_causal_lm_loss` is enabled (must be >= 1) | Lower (e.g. 256--512) for more memory savings per chunk |
+| `training.memory.debug_loss_memory` | `false` | Log coarse CUDA memory buckets around the causal-LM fp32 logits upcast for diagnostics. Logging only -- does not change training behavior. DP path only. Treat as internal debugging output, not a released DP artifact | Enable when diagnosing where DP training OOMs |
 | `training.num_input_records_to_sample` | `"auto"` | Records the model sees during training -- proxy for training time (`"auto"` or int) | First knob to increase if quality is low |
 | `training.lora_r` | `32` | LoRA rank; lower values produce fewer trainable parameters | 16--64 typical; 32 is a reasonable default |
 | `training.lora_alpha_over_r` | `1.0` | LoRA scaling ratio (alpha / rank) | Leave at 1.0 |
@@ -156,6 +161,23 @@ for the full field list.
 | `training.validation_ratio` | `0.0` | Fraction of training data held out for validation loss monitoring | Leave at 0.0 unless you specifically want to monitor validation loss |
 | `training.max_vram_fraction` | `0.8` | Fraction of total GPU VRAM to allocate for training. Must be in [0, 1] | Lower if other GPU consumers are active on the same device |
 
+Pass `memory` as a nested dictionary when configuring these controls through
+the Python SDK. A canonical dotted keyword is also accepted through `**`
+expansion. Bare memory-control fields are not direct `with_train()` keyword
+arguments; an attempted bare field reports the accepted nested path.
+
+```python
+from nemo_safe_synthesizer.sdk.library_builder import SafeSynthesizer
+
+synthesizer = SafeSynthesizer().with_train(
+    memory={"chunked_causal_lm_loss": True, "chunked_causal_lm_loss_tokens": 256}
+)
+
+synthesizer = SafeSynthesizer().with_train(
+    **{"memory.chunked_causal_lm_loss": True}
+)
+```
+
 !!! note "validation_ratio vs holdout"
     `training.validation_ratio` splits the training data to monitor
     validation loss during fine-tuning. `data.holdout` splits the full
@@ -288,6 +310,7 @@ learn about any individual record. Safe Synthesizer implements DP-SGD
 | `privacy.epsilon` | `8.0` | Privacy budget -- lower values give stronger privacy | 4.0--12.0 typical; values below 4.0 may make convergence difficult |
 | `privacy.delta` | `"auto"` | Privacy failure probability (`"auto"` or float) | Leave at `"auto"` |
 | `privacy.per_sample_max_grad_norm` | `1.0` | Max L2 norm for per-sample gradients | Leave at 1.0 |
+| `privacy.grad_sample_mode` | `"hooks"` | Opacus per-sample gradient mode: `"hooks"` for standard GradSampleModule behavior, `"ghost"` for Fast/Ghost Gradient Clipping | Use `"ghost"` when DP training OOMs from per-sample gradient memory on supported layers |
 
 Compatibility constraints:
 

docs/user-guide/running.md

@@ -714,21 +714,23 @@ Based on testing, some trade-offs identified compared to SmolLM3 on average:
 ### Quantization
 
 Enabling quantization reduces VRAM consumption at the cost of some numerical
-precision. Set `training.quantize_model` to `true` and choose a bit width with
-`training.quantization_bits`.
+precision. Set `training.quantize_model` to `true` and choose a scheme with
+`training.quantization_scheme`. New configs should set the scheme directly;
+legacy configs that still set `training.quantization_bits` continue to map
+`4` to `bnb-4bit` and `8` to `bnb-8bit`.
 
 | Setting | VRAM | Precision | Speed | Notes |
 |---------|------|-----------|-------|-------|
 | No quantization | Highest | Full | Baseline | Use when VRAM is not a constraint |
-| 8-bit | ~50% reduction | Near-full | Slightly slower | Good balance for most cases |
-| 4-bit | ~75% reduction | Reduced | Faster | Use when VRAM is tight; may affect output quality |
+| `bnb-8bit` | ~50% reduction | Near-full | Slightly slower | Good balance for most cases |
+| `bnb-4bit` | ~75% reduction | Reduced | Faster | Use when VRAM is tight; may affect output quality |
 
 === "CLI"
 
     ```bash
     safe-synthesizer run \
       --training__quantize_model true \
-      --training__quantization_bits 4 \
+      --training__quantization_scheme bnb-4bit \
       --data-source data.csv
     ```
 
@@ -740,7 +742,7 @@ precision. Set `training.quantize_model` to `true` and choose a bit width with
     synthesizer = (
         SafeSynthesizer()
         .with_data_source("data.csv")
-        .with_train(quantize_model=True, quantization_bits=4)
+        .with_train(quantize_model=True, quantization_scheme="bnb-4bit")
     )
     ```
 
@@ -749,7 +751,7 @@ precision. Set `training.quantize_model` to `true` and choose a bit width with
     ```yaml
     training:
       quantize_model: true
-      quantization_bits: 4
+      quantization_scheme: bnb-4bit
     ```
 
 ### Attention Backends

docs/user-guide/troubleshooting.md

@@ -153,12 +153,29 @@ stack traces. If you see `torch.cuda.OutOfMemoryError`:
    how to lower `training.rope_scaling_factor`, truncate records, or simplify
    grouped examples. Longer sequences require more activation memory even with
    gradient checkpointing enabled
-3. Verify `training.batch_size` is `1` (the default). The effective batch
-   size is `batch_size * gradient_accumulation_steps` (default 1 x 8 = 8).
-   Peak memory is set by the forward/backward pass on one micro-batch --
+3. Verify `training.batch_size` is `1` (the default). With integer
+   `gradient_accumulation_steps`, the effective batch size is
+   `batch_size * gradient_accumulation_steps` (default 1 x 8 = 8). Peak memory
+   is set by the forward/backward pass on one micro-batch --
    `gradient_accumulation_steps` controls how many micro-batches accumulate
    before each optimizer step but does not affect peak memory
-4. Lower `training.max_vram_fraction` (default `0.8`) to leave headroom for
+4. Cap the physical microbatch with `training.max_physical_batch_size` if you
+   need to preserve a larger logical batch. For example, `batch_size: 8`,
+   `gradient_accumulation_steps: auto`, and `max_physical_batch_size: 4`
+   trains with physical microbatches of 4 and two accumulation steps
+5. For DP runs, set `privacy.grad_sample_mode: ghost` to use Opacus
+   Fast/Ghost Gradient Clipping, which reduces per-sample gradient memory for
+   supported layers
+6. For DP runs that OOM at the loss step on large-vocab models, set
+   `training.memory.chunked_causal_lm_loss: true` (tune
+   `training.memory.chunked_causal_lm_loss_tokens`, default `1024`) to compute
+   cross entropy in token chunks instead of upcasting all logits to fp32 at
+   once, and/or `training.memory.disable_dp_bf16: true` to keep model outputs
+   in their original dtype through the loss. Set
+   `training.memory.debug_loss_memory: true` to log coarse CUDA memory buckets
+   around the logits upcast while diagnosing. Keep those diagnostics internal;
+   do not release them as part of a claimed-DP artifact
+7. Lower `training.max_vram_fraction` (default `0.8`) to leave headroom for
    other GPU consumers on the same device
 
 GPU memory during LoRA SFT breaks down into three components:
@@ -419,24 +436,33 @@ Several defaults may not match your expectations:
 
 | Parameter | Default | Notes |
 |-----------|---------|-------|
-| `training.batch_size` | `1` | Effective batch = `batch_size` x `gradient_accumulation_steps` (8) |
+| `training.batch_size` | `1` | Effective batch = `batch_size` x `gradient_accumulation_steps` (8) when accumulation is an integer |
+| `training.max_physical_batch_size` | `"auto"` | Leaves physical batching unchanged unless set to a numeric cap |
 | `training.validation_ratio` | `0.0` | No validation split by default |
 | `data.holdout` | `0.05` | 5% of records held out for evaluation; capped by `data.max_holdout` (2000) |
 | `data.random_state` | `None` | Auto-generates a random seed -- set this value explicitly if you need reproducibility |
 | `generation.num_records` | `1000` | May be too small for production use |
 
 ### Auto-Resolved Parameters
 
-Many parameters accept `"auto"` and are resolved at runtime by the
-[`AutoConfigResolver`][nemo_safe_synthesizer.config.autoconfig.AutoConfigResolver].
-See [Configuration Reference](configuration.md) for the full list.
+Many parameters accept `"auto"` and are resolved at runtime. Most are resolved
+by the
+[`AutoConfigResolver`][nemo_safe_synthesizer.config.autoconfig.AutoConfigResolver];
+training batching values resolve when Trainer arguments are built. See
+[Configuration Reference](configuration.md) for the full list.
 
 - `training.rope_scaling_factor` -- auto-estimated from dataset token counts;
   see [Context Length and Record Fitting](#context-length-and-record-fitting)
   for details and caveats
 - `training.num_input_records_to_sample` -- derived from `rope_scaling_factor * 25000`
 - `training.learning_rate` -- model-specific default from `ModelMetadata`:
   Mistral uses 0.0001, all other supported model families use 0.0005
+- `training.max_physical_batch_size` -- leaves physical batching unchanged
+  when `"auto"`; a numeric value caps the Trainer microbatch and preserves the
+  logical batch by increasing accumulation when possible
+- `training.gradient_accumulation_steps` -- derives accumulation from
+  `training.batch_size` and `training.max_physical_batch_size` when set to
+  `"auto"`
 - `data.max_sequences_per_example` -- resolves to `1` when differential
   privacy is enabled (required to limit per-example gradient contribution),
   `10` otherwise for best performance

pyproject.toml

@@ -118,7 +118,7 @@ cpu = [
   "gliner",
   "kernels>=0.12.1,<0.15",
   "peft>=0.18.0",
-  "opacus",
+  "opacus>=1.6.0",
   "sentence-transformers",
   "torch==2.11.0; sys_platform == 'darwin'",
   "torch==2.11.0+cpu; sys_platform == 'linux'",
@@ -144,7 +144,7 @@ cu129 = [
   "kernels>=0.12.1,<0.15",
   "nvidia-cublas-cu12; sys_platform == 'linux'",
   "nvidia-ml-py; sys_platform == 'linux'",
-  "opacus",
+  "opacus>=1.6.0",
   "peft>=0.18.0",
   "sentence-transformers",
   "torch==2.11.0+cu129; sys_platform == 'linux'",

src/nemo_safe_synthesizer/config/autoconfig.py

@@ -308,9 +308,9 @@ def _build_updated_params(
             "data": data_params,
             "privacy": privacy_params,
         }
-        logger.debug(f"params to update: {new_params}")
+        logger.debug("params to update", extra={"ctx": new_params})
         my_config = self._config.with_config_patch(new_params)
-        logger.debug(f"auto-updated config: {my_config.model_dump(exclude_unset=True)}")
+        logger.debug("auto-updated config", extra={"ctx": my_config.model_dump(exclude_unset=True)})
         return my_config
 
     def resolve(self) -> SafeSynthesizerParameters:

src/nemo_safe_synthesizer/config/differential_privacy.py

@@ -5,6 +5,7 @@
 
 from typing import (
     Annotated,
+    Literal,
 )
 
 from pydantic import (
@@ -73,3 +74,16 @@ class DifferentialPrivacyHyperparams(Parameters):
             description="Maximum L2 norm for per-sample gradient clipping. Must be > 0.",
         ),
     ] = 1.0
+
+    grad_sample_mode: Annotated[
+        Literal["hooks", "ghost"],
+        Field(
+            title="grad_sample_mode",
+            description=(
+                "Opacus per-sample gradient mode for DP training. "
+                "'hooks' uses standard Opacus GradSampleModule behavior. "
+                "'ghost' uses Opacus Fast/Ghost Gradient Clipping to reduce "
+                "per-sample gradient memory for supported layers."
+            ),
+        ),
+    ] = "hooks"