TrainLoraNode: offloading can silently change forward/recompute values during gradient checkpointing on MPS

### Custom Node Testing

- [x] I have tried disabling custom nodes and the issue persists (see [how to disable custom nodes](https://docs.comfy.org/troubleshooting/custom-node-issues#step-1%3A-test-with-all-custom-nodes-disabled) if you need help)

### Expected Behavior

If a module is recomputed during backward, offloading should not cause the
recomputed numerical result to differ materially from the result used in the
original forward. Recomputation is supposed to reproduce the forward, not
merely produce a tensor of the same shape.

### Actual Behavior

With `gradient_checkpointing=true` and `offloading=true`, I measured a large
numerical difference between activations produced during the original forward
and those produced when the checkpointed computation is recomputed during
backward.

At 0.25 MP this does **not** raise a `torch.utils.checkpoint.CheckpointError`.
Tensor shapes, dtypes and devices remain unchanged while tensor values differ,
and the stock `determinism_check` compares only shape/dtype/device — so the
training run completes and reports success.

Paired measurement, same configuration, only `offloading` changed:

```
offloading=true:
sampled activation std:
forward 0.208552
recomputation 0.257003
relative difference ~19%

offloading=false:
no sampled-value divergence observed
maximum recorded relative difference = 0.0
```

The `offloading=false` control is why I do not rely on any assumed numerical
noise threshold for this claim: the same instrumentation, on the same pipeline,
reports exactly `0.0` when offloading is disabled.

**Impact.** The backward recomputation is not numerically reproducing the
forward computation, so the resulting gradients need not correspond to the
original forward evaluation. This is particularly concerning because the run
otherwise appears entirely successful: no OOM, no NaN, no exception, nothing in
the logs.

**Relevant implementation.** `nodes_train.py` contains two different
checkpointing mechanisms operating on the same model. For Linear/Conv modules
when offloading is enabled, `OffloadCheckpointFunction` (line ~633) executes
the forward under `torch.no_grad()`, stores the input and callable, and
executes the callable again in backward. Other patched modules are wrapped in
`torch.utils.checkpoint.checkpoint(..., use_reentrant=False)` (line ~685).
The divergence reported here is associated specifically with the
`offloading=true` path.

### Steps to Reproduce

`TrainLoraNode` on Z-Image Base bf16, MPS:

| setting | value |
|---|---|
| resolution | 0.25 MP |
| `gradient_checkpointing` | `true` |
| `checkpoint_depth` | 2 |
| `offloading` | `true` (then `false` for the control) |
| batch size | 1 |
| `training_dtype` | bf16 |

Run the same training twice, changing only `offloading`.

With stock ComfyUI both runs complete without error, which is precisely the
problem — the difference is not observable without instrumentation.

**Instrumentation used for diagnosis only.** The bug exists on stock ComfyUI;
a local build was used only to make it measurable. It replaced the checkpoint
determinism-check metadata extractor so that it also compares values. To avoid
materially changing memory pressure — the very variable the suspected mechanism
is sensitive to — the final instrumentation does not copy whole activation
tensors to float32. For each compared tensor it uses:

- a deterministic sample of approximately 4096 elements at fixed stride;
- 64 additional probes at fixed deterministic positions;
- exactly the same sample positions for forward and recomputation.

It records statistics only and does not alter the training computation.

A methodological note for anyone reproducing this: an earlier version of the
instrumentation computed statistics over whole tensors via a float32 copy,
which itself raised memory pressure. That measurement should not be trusted.
Also, relative differences on `mean` are unreliable here because the mean is
near zero and subject to cancellation; `std` is the more stable witness.

### Debug Logs

```powershell
With `offloading=true`, one representative comparison at 0.25 MP:

Gradient checkpointing: patching 40 modules at depth 2
...
[instrumented determinism check] saved vs recomputed
std 0.208552 vs 0.257003 relative difference 1.885e-01
std 25.736 vs 25.736 relative difference 2.2e-07 (healthy reference)

With `offloading=false`, identical configuration and identical instrumentation:

Gradient checkpointing: patching 40 modules at depth 2
...
[instrumented determinism check] maximum relative difference over all
compared tensors: 0.0

Both runs logged `Gradient checkpointing: patching 40 modules at depth 2`, so
the comparisons were actually performed in both cases.
```

### Other

**Environment**

- ComfyUI commit `4da9e2dbead52fc1e68beae33fe3d7ad63b63241`, base tag `v0.33.3`
- `comfy_extras/nodes_train.py` as of commit `d0fec2ef` (2026-07-21)
- PyTorch 2.10.0, Python 3.13.12
- macOS 26.6.2, Apple Silicon M5 Pro, 24 GB unified memory, device MPS
- Model: Z-Image Base bf16, batch size 1, 0.25 MP

**Workaround**

`offloading=false` removes the measured divergence. On this 24 GB machine the
cost is negligible (~5.1-5.2 s/step versus ~5.5 s/step with offloading
enabled), and with PyTorch attention and `checkpoint_depth=2` the model still
trains without OOM.

**Possible directions for a fix** — I have not implemented one, so these are
directions rather than a proposed patch:

- avoid nesting/interleaving `OffloadCheckpointFunction` with PyTorch's
non-reentrant checkpoint on the same model;
- have the offload path preserve whatever state is required to make its
recomputation deterministic;
- optionally, extend the checkpoint determinism check with an opt-in
value-comparison mode, so this class of divergence is detectable rather than
silent.

**Possibly related**

- #10940 — `TrainLoraNode` device mismatch with offloading enabled. Different
symptom and platform (CUDA), same offload code path.

Start in comfy_extras/nodes_train.py at OffloadCheckpointFunction around line 633 and the non-reentrant checkpoint wrapping around line 685. Reproduce the paired MPS configuration with gradient_checkpointing and offloading enabled and disabled, then trace how forward inputs, callable state, and recomputation interact. Done means the offloading path no longer silently produces materially different recomputed values, or the divergence is reliably detected.

Written by the indexing model from the issue text.