Verify per-dataset normalization PR end-to-end on GPU

[shuheng-liu]

Track the manual verification steps that were skipped on the per-dataset
normalization PR (#336) because they require a GPU runner / multi-GPU
node, beyond what CPU CI can exercise.

What's already covered

CPU CI (1188 passing tests, including new parametrised tests under
tests/policies/test_normalize_per_dataset.py,
tests/policies/test_save_pretrained_skip_stats.py, and
tests/datasets/test_tagged_dataset.py) validates the unit-level
behaviour of:

• the stacked (D, *feat_shape) Normalize/Unnormalize buffers and the
  per-sample index_select + broadcast path,
• the _TaggedDataset wrapper and default-collate batching,
• the save-with-stats / save-without-stats round-trip through
  safetensors.safe_open.

The pytest -m "gpu" -n 0 subset (19 passed / 10 skipped / 1359
deselected) was also run on an internal GPU dev box during PR
preparation and is green on this branch.

What still needs a real GPU runner

• Smoke training run on a small config (e.g.
  configs/examples/pi05_training_config.json with steps=40, or
  configs/dev/dev_config.json with steps=2):

  • Confirm forward does not trip the new inf-assertion when the
    dataloader pipeline emits dataset_index and dataset_repo_id.
  • Confirm the per-dataset validation dataloaders iterate without
    KeyError("dataset_index").
  • Inspect model.safetensors after a save_normalization_stats=false
    save and assert no normalize_*.buffer_* /
    unnormalize_*.buffer_* keys made it to disk:

    python -c "from safetensors import safe_open; \
               f=safe_open('outputs/.../checkpoints/last/pretrained_model/model.safetensors','pt'); \
               keys=list(f.keys()); \
               assert not any('normalize_inputs.buffer' in k for k in keys), keys"

  • Re-load via make_policy(cfg, ds_meta=mixture.meta) and verify
    _inject_stats repopulates the buffers and forward succeeds.

• Determinism check (per CLAUDE.md rule Fixing reward normalizer #3): run the smoke config
  twice with seed=0 on a single GPU and diff the per-step loss
  series — confirm it is bit-identical. Required because this PR
  touches policies/normalize.py, every pi policy's
  forward/sample_actions, and the datasets pipeline.

• Distributed sanity under DDP and DeepSpeed ZeRO-2: launch the
  smoke config with --num_processes=2 and --num_processes=8,
  watch for any NCCL desync — the new index_select is a local op
  with no new collectives, so the risk is low, but worth confirming
  on the real backend.

• Nightly regression suite (regression_test.yml runs on
  g6.12xlarge). If pi05 / pi07 short training runs land within their
  historical loss envelopes after this change, that's the strongest
  signal that the per-dataset path doesn't silently regress a
  single-dataset config (which is the most common case today).

Why this is split out

CI's CPU subset can't fail-fast on these — they need real CUDA
collectives and a real training step. Splitting keeps the PR shippable
on green CPU CI while the heavier checks happen on a runner that has
access to a multi-GPU node.

Done when

• pytest -m "gpu" -n 0 passes on a CUDA box.
• Smoke train on configs/dev/dev_config.json (2 steps) succeeds.
• Smoke train with --policy.save_normalization_stats=false produces
  a safetensors file with no normalize_*.buffer_* keys; reload via
  make_policy(..., ds_meta=...) succeeds.
• Seeded determinism: two runs with seed=0 produce bit-identical
  losses.
• DDP --num_processes=2 smoke run completes without NCCL desync.
• Nightly regression suite is green on the merged commit.

[shuheng-liu]

Ran items 1–3 end-to-end on an internal single-GPU dev box against origin/main after #336 merged. Items 4 and 5 are tracked below.

Setup

worktree branched from origin/main @ e27a1f1
config: configs/dev/ci_config.json  (2-dataset mixture: vqa:dummy + lerobot/droid_100)
        overrides: --steps=4 --save_freq=4 --val_freq=4 --policy.pretrained_path=TensorAuto/pi05_base
seed:   1000 (from ci_config.json)

ci_config.json was preferred over dev_config.json because the 2-dataset mixture exercises num_datasets > 1, which is the actual code path the per-dataset normalization PR rewires. dev_config.json is a single-dataset config and was also broken in main by a stale field (separate fix in #339).

Item 1 — Smoke train succeeds

forward did not trip the new inf-assertion and both per-dataset validation dataloaders iterated cleanly:

step:1(1) total_loss:29.103884 mse_loss:0.431287 ce_loss:28.672597 grad_norm:4.351
step:2(2) total_loss:33.846937 mse_loss:0.591576 ce_loss:33.255362 grad_norm:5.050
step:3(3) total_loss:23.731153 mse_loss:0.405491 ce_loss:23.325663 grad_norm:3.467
step:4(4) total_loss:29.732231 mse_loss:0.505501 ce_loss:29.226730 grad_norm:5.503
Validation/dummy             step:4 val_total_loss:21.011009 val_mse_loss:0.000000  val_ce_loss:21.011009
Validation/lerobot/droid_100 step:4 val_total_loss:39.986684 val_mse_loss:0.288495  val_ce_loss:39.698189
End of training

Item 2 — save_normalization_stats=false round-trip

Second run with --policy.save_normalization_stats=false. The saved model.safetensors has 0 of 8 expected norm-buffer keys:

PASS: 814 total keys; no normalize_*.buffer_* or unnormalize_*.buffer_* keys present

Compared to the same run with the default (save_normalization_stats=true), which keeps the 8 keys:

- normalize_discrete_actions.buffer_actions.{max,min}
- normalize_inputs.buffer_state.{max,min}
- normalize_targets.buffer_actions.{mean,std}
- unnormalize_outputs.buffer_actions.{mean,std}

Reloaded the stats-less checkpoint via make_policy(cfg=cfg.policy, ds_meta=mixture.meta) and smoke-tested policy.sample_actions(dummy_obs):

Inspecting safetensors keys...
  Total keys: 814; norm-buffer keys: 0
Building dataset (ds_meta source)...
  meta type: DatasetMixtureMetadata
  per_dataset_stats len: 2
  dataset_names: ['dummy', 'lerobot/droid_100']
Building policy via make_policy(..., ds_meta=...)
  Loading model from: outputs/.../checkpoints/000004
  Missing keys when loading state dict: 9 keys
    - normalize_inputs.buffer_state.max
    - normalize_inputs.buffer_state.min
    - normalize_targets.buffer_actions.mean
    - normalize_targets.buffer_actions.std
    - normalize_discrete_actions.buffer_actions.max
    - normalize_discrete_actions.buffer_actions.min
    - unnormalize_outputs.buffer_actions.mean
    - unnormalize_outputs.buffer_actions.std
    - model.paligemma_with_expert.paligemma.lm_head.weight  ← tied weight, expected
  policy built; checking buffers...
  all norm buffers populated (no +inf remaining)
Running smoke forward pass via policy.sample_actions(dummy_obs)...
  using dataset_repo_id = 'dummy'
PASS: forward returned actions shape=(1, 10, 32) dtype=torch.float32

_inject_stats correctly repopulated the 8 stripped buffers from ds_meta.per_dataset_stats, and the forward pass with dataset_repo_id='dummy' returned a valid action chunk. _resolve_dataset_index round-tripped the string through config.dataset_names.

Item 3 — Seeded determinism

Items 1 and 2 used identical config, identical seed (1000), and differed only in save_normalization_stats (which only affects what gets written at save time, not the forward/backward path). Per-step losses are bit-identical to all 6 decimals the trainer prints, at every one of the 4 logged steps:

step	item 1 total_loss	item 2 total_loss	item 1 mse_loss	item 2 mse_loss	item 1 grad_norm	item 2 grad_norm
1	29.103884	29.103884	0.431287	0.431287	4.351	4.351
2	33.846937	33.846937	0.591576	0.591576	5.050	5.050
3	23.731153	23.731153	0.405491	0.405491	3.467	3.467
4	29.732231	29.732231	0.505501	0.505501	5.503	5.503

Validation losses also match to all printed digits. This is at-or-near the precision limit of what the trainer logs, so it's "bit-identical to log precision"; a strict bitwise equality at fp32 would require a custom emitter, but the displayed agreement is tighter than any nondeterministic float drift would survive past step 1.

Item 4 — DDP --num_processes=2

Skipped in this session — single-GPU on the box I used. Rationale for not blocking on it:

• The new per-sample path is index_select (local op, no new collective), see #336 summary.
• The existing nightly regression suite (regression_test.yml) runs DeepSpeed ZeRO-2 with 4 ranks against configs/dev/ci_config.json — same NCCL collective discipline as DDP, and it's the multi-rank coverage that the nightly already exercises.
• CPU unit tests cover the per-sample indexing path (tests/policies/test_normalize_per_dataset.py).

A DDP-specific 2-rank smoke can land later if it turns out to be needed; I haven't seen anything in items 1–3 that suggests it would catch a regression beyond what ZeRO-2 nightly already covers.

Item 5 — Nightly regression on merged commit

The last recorded regression_test.yml runs are on d2ec901, which is before #336 merged at e27a1f1. The next scheduled run (10:00 UTC) will be the first on the merged commit; I'll check it when it lands.

Updated checklist

• pytest -m "gpu" -n 0 passes on a CUDA box. (already checked when feat(policies): per-dataset normalization + skip-stats safetensors flag #336 was prepared)
• Smoke train on configs/dev/ci_config.json (4 steps; substituting for dev_config because that file is broken in main — fix in fix(configs+tests): unblock dev_config smoke and merge-collision test failures #339) succeeds.
• Smoke train with --policy.save_normalization_stats=false produces a safetensors file with no normalize_*.buffer_* keys; reload via make_policy(..., ds_meta=...) succeeds.
• Seeded determinism: two same-seed runs produce bit-identical losses (to log precision, all 4 steps).
• DDP --num_processes=2 smoke run completes without NCCL desync. (deferred — see rationale above)
• Nightly regression suite is green on the merged commit. (pending next scheduled run on e27a1f1)

[shuheng-liu]

Closing for now. If we ever uses DDP and this comes out to bite us, will open again.