TrainLoraNode: memory-dependent attention chunking can change between forward and gradient-checkpoint recomputation

### 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

Gradient-checkpoint recomputation should execute a computation structurally
compatible with the original forward, independently of transient memory
pressure. Training should not fail depending on how much memory happened to be
free at the moment each attention call was made.

### Actual Behavior

`TrainLoraNode` fails during backward with
`torch.utils.checkpoint.CheckpointError` because the attention implementation
chooses its chunk size from the amount of free memory available at runtime.
The original forward and the gradient-checkpoint recomputation can observe
different free-memory values and therefore split the same attention sequence
differently.

**Root cause, isolated by a paired A/B run** (details under Steps to Reproduce):
in `comfy/ldm/modules/attention.py`, the sub-quadratic attention path calls
`model_management.get_free_memory(...)` and then selects `query_chunk_size`
from the list:

```
[4096, 2048, 1024, 512, 256]
```

by finding the largest `x` for which
`mem_free_total / (batch_x_heads * bytes_per_token * x * 4.0) >= k_tokens`.

For the failing call in my reproduction:

- `batch_x_heads = 30`
- dtype bf16, so `bytes_per_token = 2`
- sequence length 2176 tokens

Under that heuristic, a 4096 chunk requires roughly 2.14 GB free and a 2048
chunk roughly 1.07 GB. The forward and the recomputation land on opposite sides
of that threshold:

```
forward: query_chunk_size = 4096 -> one effective 2176-token chunk
recomputation: query_chunk_size = 2048 -> 2048 + 128
```

so the first differing saved tensor is `[30, 2176, 128] -> [30, 2048, 128]`,
and every subsequent saved tensor is displaced in the checkpoint sequence.
That displacement is why the exception reports dozens of downstream metadata
mismatches rather than a single one.

Note that `2048` is not an arbitrary number: it is literally one of the
constants in the list above.

This makes runtime free memory part of the computational structure inside a
function that is later recomputed by gradient checkpointing.

**Second isolated variant.** With `--use-split-cross-attention` I observe a
different deterministic mismatch on the same 2176-token sequence, which becomes
1088 during recomputation. That attention path also makes its splitting
decision dynamically from available memory.

I consider the memory-dependent attention splitting root cause isolated for
these two variants only. I have observed other `CheckpointError` variants on
this node which I have **not** isolated, and I am not claiming they share this
cause.

### Steps to Reproduce

Two runs, identical except for the attention backend.

**Common configuration** (Z-Image Base bf16, MPS):

| setting | value |
|---|---|
| dataset | 32 image/text pairs |
| `ImageScaleToTotalPixels` | `megapixels=0.5`, `resolution_steps=16` |
| `TrainLoraNode` | `steps=1`, `seed=0`, `batch_size=1`, rank 16, AdamW, MSE |
| | `gradient_checkpointing=true`, `checkpoint_depth=2` |
| | `offloading=false`, `bucket_mode=true`, `training_dtype=bf16` |

Resolution buckets produced: `110x74:15`, `74x110:11`, `124x66:1`,
`100x82:1`, `154x52:1`, `84x98:1`, `96x84:1`, `118x70:1`.
Both runs deterministically selected bucket 3.

**Run A — default attention.** Launch args `--enable-manager`.
Startup log line: `Using sub quadratic optimization for attention`.
Result: **fails** at step 0 inside `bwd_loss.backward()`.

**Run B — PyTorch attention.** Launch args
`--enable-manager --use-pytorch-cross-attention`.
Startup log line: `Using pytorch attention`.
Result: **passes**, forward and backward both complete.

The two runs were performed seven minutes apart on the same machine and
session. `offloading` was `false` in both, so the custom
`OffloadCheckpointFunction` path in `nodes_train.py` is not required to trigger
this failure.

There was ~14.2 GB reported free on MPS when the exception was raised, so this
is not an out-of-memory failure.

**On reproducer size:** I have not reduced the dataset below 32 images. The
failing branch explicitly depends on runtime free memory, so shrinking the
dataset may change the threshold crossing itself and make the failure disappear
for reasons unrelated to the defect. The image content is irrelevant; the
memory state produced by the pipeline is not.

### Debug Logs

```powershell
Run A (fails):

Using sub quadratic optimization for attention
...
Loaded 32 images from .../input/
Resolution bucket (110x74): 15 samples
Resolution bucket (74x110): 11 samples
Resolution bucket (124x66): 1 samples
Resolution bucket (100x82): 1 samples
Resolution bucket (154x52): 1 samples
Resolution bucket (84x98): 1 samples
Resolution bucket (96x84): 1 samples
Resolution bucket (118x70): 1 samples
Created 8 resolution buckets from 32 samples
model weight dtype torch.bfloat16, manual cast: None
model_type FLOW
Gradient checkpointing: patching 40 modules at depth 2
Requested to load Lumina2
loaded completely; 11825.14 MB loaded, full load: True
Training LoRA: 0%| | 0/1 [00:02

Start in comfy/ldm/modules/attention.py by reading the sub-quadratic and split cross-attention paths and their free-memory-based chunking decisions. Reproduce with gradient_checkpointing enabled and the default attention backend, then inspect comfy_extras/nodes_train.py around fwd_bwd and checkpointing. Done means forward and recomputation use compatible attention structure and the reported CheckpointError no longer occurs.

Written by the indexing model from the issue text.