Use fused attention, make the mask fill fp16-safe, validate the AMP dtype, expose DataLoader worker options

[Rian354]

Summary

Performance and correctness in the attention path and the data loader.

The explicit attention calculation builds several (B, H, S, S) tensors for each layer. Only an
interpretability pass reads them. Ordinary training now uses fused SDPA, and the explicit path
stays available behind register_hook=True.

Two silent defects are also repaired. The attention mask fill value -1e9 is outside the fp16
range, so the explicit path is not safe under fp16 mixed precision. The mixed-precision dtype was
not validated, so any value other than "bf16" selected fp16 with no message.

This PR does not add mixed precision. main already has use_amp, amp_dtype and correct
GradScaler gating for fp16.

---

Implemented

1. Fused attention with scaled_dot_product_attention

• File: pyhealth/models/transformer.py
• MultiHeadedAttention.forward now uses F.scaled_dot_product_attention when the caller does
  not request the attention map.
• The explicit path remains for register_hook=True, because an interpretability pass needs the
  map and its gradient.
• The two paths agree. The maximum absolute difference is 7.153e-07 with no padding and
  4.768e-07 with padding.
• The fused path also clears attn_map and attn_gradients. Before this change the map stayed
  in memory after every forward pass, and no caller read it.

2. Mask fill value that is safe in half precision

• File: pyhealth/models/transformer.py
• The explicit path used scores.masked_fill(pad_mask, -1e9). The value -1e9 is outside the
  range of fp16, so the calculation is not correct under fp16 mixed precision.
• The fill value now comes from the dtype: torch.finfo(scores.dtype).min.
• A second masked_fill after the softmax sets masked positions to exactly 0.

3. Validation for the mixed-precision dtype

• File: pyhealth/trainer.py
• The previous expression was torch.bfloat16 if amp_dtype == "bf16" else torch.float16. Every
  other value therefore selected fp16 with no message. A typo, or the spelling "bfloat16",
  changed the precision of a run in silence.
• fp16 also needs a GradScaler, so the silent path changed the gradient behaviour as well.
• Added resolve_amp_dtype(). It accepts bf16, bfloat16, fp16 and float16. It raises for
  any other value. It also raises when bf16 is requested on a CUDA device without bf16 support.

4. DataLoader worker options

• File: pyhealth/datasets/utils.py
• Added num_workers, pin_memory, persistent_workers and prefetch_factor to
  get_dataloader.
• An invalid combination raises. persistent_workers needs num_workers above 0.
  prefetch_factor needs num_workers above 0 and a positive value.
• litdata.StreamingDataLoader is selected only when num_workers is above 0. It coordinates
  shard reads between workers, and it gives no advantage in a single process.
• The default behaviour does not change.

---

Validation

The two attention paths agree

Condition	Maximum absolute difference
All positions valid	7.153e-07
With padding	4.768e-07

Measured effect of mixed precision

These numbers characterise the behaviour that main already has. No measurement of it exists
today, and the projected cost of pretraining depends on it.

A10 GPU, notes_labs, transformer 128/2/4, full scale, 2 epochs, batch 8:

Precision	Wall time	Seconds for each epoch	Peak VRAM	train_loss
bf16	5,275 s	1471.3, 1007.2	1,814 MB	1.2345 -> 1.1412
fp32	10,198 s	4176.1, 1808.8	2,402 MB	1.2348 -> 1.1341

Speed for each epoch is 2.84x for epoch 1 and 2.41x as a mean of the two epochs. Peak
VRAM decreases by 24.5%. The final training losses agree to 0.63%.

An earlier measurement on a --dev 2500 subset gave 3.6x to 3.8x. On that subset the fixed
startup cost is a large part of each epoch. The full-scale numbers above are the correct values.

The DataLoader options give no advantage here

Measured with mixed precision already enabled: 39.0 and 41.8 seconds for each epoch by default,
against 40.0 and 43.0 seconds when tuned. There is no advantage. The options are available and
are not enabled. This result is reported so that a reader does not repeat the measurement.

---

Unit tests

File: tests/test_attention_paths.py, 7 tests.

Test	Purpose
test_fused_and_explicit_paths_agree_without_padding	Regression guard
test_fused_and_explicit_paths_agree_with_padding	Padding shows a mask convention error
test_fused_path_does_not_retain_the_attention_map	The map is not kept when no caller reads it
test_explicit_path_still_supplies_the_attention_map	Interpretability still works
test_mask_fill_value_is_representable_in_half_precision	Fails on main
test_amp_dtype_is_validated_not_silently_coerced	Fails on main
test_gradients_flow_through_the_fused_path	Regression guard

Two of the seven tests fail on main. These are the two defect tests. The other five are
regression guards: main has one attention path, so agreement between two paths is true there
without any change.

No regression: 12 existing transformer and dataset tests pass on this branch.

---

Not in this PR

The Optuna changes are held for a later PR. scripts/optuna_e2e.py imports
fit_lab_standardizer and lab_standardizer_fit_scope, and scripts/optuna_pretrain.py imports
pyhealth.models.pretrain. main has none of these, so an Optuna PR cannot import until the
pretraining PR and the observation-window PR are merged.

Draft #43 has the same problem today. scripts/optuna_pretrain.py:65 imports
pyhealth.models.pretrain, and that package is not on the #43 branch.

The cache for the frozen text encoder is also held for the pretraining PR. The cache is selected
by field_name in self._frozen_text_fields, and main has no _frozen_text_fields and no
text_finetune_mode.

Merge order

This PR shares no file with the observation-window PR (#46), so the two can merge in either
order.

It does share files with the later PRs: pyhealth/models/transformer.py and
pyhealth/datasets/utils.py with #48, and pyhealth/trainer.py with #49. Merge #46, then #48,
then #49, and rebase this PR onto the result. The changes do not depend on each other, but a
clean merge at any time will conflict.

Questions for the reviewer

1. The fused path returns no attention map. Does any caller outside the interpretability path
   read attn_map today?
2. resolve_amp_dtype now raises for a value that previously selected fp16. Is any configuration
   in use that depends on the previous behaviour?