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

pyhealth/datasets/utils.py

@@ -327,26 +327,61 @@ def collate_fn_dict_with_padding(batch: List[dict]) -> dict:
 
 
 def get_dataloader(
-    dataset: litdata.StreamingDataset, batch_size: int, shuffle: bool = False
+    dataset: litdata.StreamingDataset,
+    batch_size: int,
+    shuffle: bool = False,
+    num_workers: int = 0,
+    pin_memory: bool = False,
+    persistent_workers: bool = False,
+    prefetch_factor: Optional[int] = None,
 ) -> DataLoader:
     """Creates a DataLoader for a given dataset.
 
     Args:
         dataset: The dataset to load data from.
         batch_size: The number of samples per batch.
         shuffle: Whether to shuffle the data at every epoch.
+        num_workers: Number of worker processes that load and collate batches.
+        pin_memory: Copy CPU tensors into page-locked memory before return.
+        persistent_workers: Keep the loader workers between epochs. This is valid
+            only when ``num_workers`` is more than 0.
+        prefetch_factor: Batches that each worker loads in advance. This is valid
+            only when ``num_workers`` is more than 0. ``None`` keeps the default
+            of PyTorch.
 
     Returns:
         A DataLoader instance for the dataset.
     """
     dataset.set_shuffle(shuffle)
-    dataloader = DataLoader(
-        dataset,
-        batch_size=batch_size,
-        collate_fn=collate_fn_dict_with_padding,
-    )
+    if num_workers < 0:
+        raise ValueError(f"num_workers must be non-negative, got {num_workers}.")
+    if persistent_workers and num_workers == 0:
+        raise ValueError("persistent_workers requires num_workers > 0.")
+    if prefetch_factor is not None and (num_workers == 0 or prefetch_factor <= 0):
+        raise ValueError(
+            "prefetch_factor must be positive and requires num_workers > 0."
+        )
 
-    return dataloader
+    loader_kwargs = {
+        "dataset": dataset,
+        "batch_size": batch_size,
+        "collate_fn": collate_fn_dict_with_padding,
+        "num_workers": num_workers,
+        "pin_memory": pin_memory,
+    }
+    if num_workers > 0:
+        loader_kwargs["persistent_workers"] = persistent_workers
+        if prefetch_factor is not None:
+            loader_kwargs["prefetch_factor"] = prefetch_factor
+
+    # StreamingDataLoader coordinates shard reads across workers. With a single
+    # process it adds no advantage, so keep the plain DataLoader there.
+    loader_class = (
+        litdata.StreamingDataLoader
+        if isinstance(dataset, litdata.StreamingDataset) and num_workers > 0
+        else DataLoader
+    )
+    return loader_class(**loader_kwargs)
 
 
 def save_processors(sample_dataset, output_dir: str) -> Dict[str, str]:

pyhealth/models/transformer.py

@@ -7,6 +7,7 @@
 from typing import Any, Dict, Optional, Tuple, Union, cast
 
 import torch
+import torch.nn.functional as F
 from torch import nn
 
 from pyhealth.datasets import SampleDataset
@@ -55,9 +56,12 @@ def forward(
             # Use -inf so softmax produces exact zeros on padded positions,
             # avoiding a second masked_fill after softmax (saves one full
             # [B, H, S, S] boolean allocation and an extra copy).
-            pad_mask = mask == 0
-            scores = scores.masked_fill(pad_mask, -1e9)
+            # Use the dtype minimum, not -1e9. Under fp16 autocast -1e9 is
+            # outside the representable range.
+            scores = scores.masked_fill(mask == 0, torch.finfo(scores.dtype).min)
         p_attn = self.softmax(scores)
+        if mask is not None:
+            p_attn = p_attn.masked_fill(mask == 0, 0)
         if dropout is not None:
             p_attn = dropout(p_attn)
 
@@ -150,19 +154,42 @@ def forward(
         ]
 
         # 2) Apply attention on all the projected vectors in batch.
-        if mask is not None:
-            mask = mask.unsqueeze(1)
-        x, attn = self.attention(query, key, value, mask=mask, dropout=self.dropout)
-
-        if register_hook:
-            # Only store attn_map and hook during interpretability passes.
-            # Using .detach() gives an independent copy whose storage
-            # is NOT shared with the live graph, so the graph can be freed
-            # normally after .backward() without leaking GPU memory.
+        # The explicit path materialises several (B, H, S, S) tensors for each
+        # layer. It is necessary only for interpretability, where a caller reads
+        # the attention map and its gradient. Ordinary training uses fused SDPA.
+        if not register_hook:
+            query_mask = None
+            attn_mask = None
+            if mask is not None:
+                valid = mask.bool()
+                if mask.dim() == 2:
+                    query_mask = valid[:, None, :, None]
+                    attn_mask = valid[:, None, None, :]
+                else:
+                    query_mask = valid.any(dim=-1)[:, None, :, None]
+                    attn_mask = valid.unsqueeze(1)
+            x = F.scaled_dot_product_attention(
+                query,
+                key,
+                value,
+                attn_mask=attn_mask,
+                dropout_p=self.dropout.p if self.training else 0.0,
+            )
+            if query_mask is not None:
+                x = x * query_mask.to(x.dtype)
+            self.attn_map = None
+            self.attn_gradients = None
+        else:
+            if mask is not None:
+                mask = mask.unsqueeze(1)
+            x, attn = self.attention(
+                query, key, value, mask=mask, dropout=self.dropout
+            )
+            # .detach() gives an independent copy whose storage is NOT shared
+            # with the live graph, so the graph can be freed after .backward()
+            # without a memory leak.
             self.attn_map = attn.detach()
             attn.register_hook(self.save_attn_grad)
-        else:
-            self.attn_map = None
         # 3) "Concat" using a view and apply a final linear.
         x = x.transpose(1, 2).contiguous().view(batch_size, -1, self.h * self.d_k)
 

pyhealth/trainer.py

@@ -20,6 +20,41 @@
 logger = logging.getLogger(__name__)
 
 
+_AMP_DTYPES = {
+    "bf16": torch.bfloat16,
+    "bfloat16": torch.bfloat16,
+    "fp16": torch.float16,
+    "float16": torch.float16,
+}
+
+
+def resolve_amp_dtype(amp_dtype: str, use_amp: bool = False) -> torch.dtype:
+    """Validate the mixed-precision dtype and return it.
+
+    The previous expression was ``bfloat16 if amp_dtype == "bf16" else float16``,
+    so every other value selected fp16 without a message. A typo, or the spelling
+    ``"bfloat16"``, silently changed the precision of a run. fp16 also needs a
+    GradScaler, so the silent path changed the gradient behaviour as well.
+    """
+    name = str(amp_dtype).lower()
+    if name not in _AMP_DTYPES:
+        raise ValueError(
+            f"amp_dtype must be one of {sorted(_AMP_DTYPES)}, got {amp_dtype!r}."
+        )
+    dtype = _AMP_DTYPES[name]
+    if (
+        use_amp
+        and dtype == torch.bfloat16
+        and torch.cuda.is_available()
+        and not torch.cuda.is_bf16_supported()
+    ):
+        raise RuntimeError(
+            "bf16 mixed precision was requested, but this CUDA device does not "
+            "support bf16."
+        )
+    return dtype
+
+
 def is_best(best_score: float, score: float, monitor_criterion: str) -> bool:
     if monitor_criterion == "max":
         return score > best_score
@@ -168,9 +203,7 @@ def train(
         if optimizer_params is None:
             optimizer_params = {"lr": 1e-3}
 
-        _amp_dtype = (
-            torch.bfloat16 if amp_dtype == "bf16" else torch.float16
-        )
+        _amp_dtype = resolve_amp_dtype(amp_dtype, use_amp)
         # GradScaler only needed for fp16; bf16 has fp32 dynamic range
         scaler = (
             torch.cuda.amp.GradScaler()

tests/test_attention_paths.py

@@ -0,0 +1,120 @@
+"""The fused and explicit attention paths must agree.
+
+The explicit path materialises several ``(B, H, S, S)`` tensors for each layer
+and exists only for interpretability. Ordinary training uses fused SDPA. If the
+two paths disagree, an interpretability pass describes a model that training
+never produced.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from pyhealth.models.transformer import TransformerLayer
+
+
+def _layer(seed: int = 0) -> TransformerLayer:
+    torch.manual_seed(seed)
+    return TransformerLayer(feature_size=128, heads=4, num_layers=2).eval()
+
+
+def test_fused_and_explicit_paths_agree_without_padding():
+    layer = _layer()
+    x = torch.randn(4, 32, 128, requires_grad=True)
+    mask = torch.ones(4, 32)
+
+    fused, _ = layer(x, mask, register_hook=False)
+    explicit, _ = layer(x, mask, register_hook=True)
+
+    assert torch.allclose(fused, explicit, atol=1e-5)
+
+
+def test_fused_and_explicit_paths_agree_with_padding():
+    """Padding is where a mask convention error shows."""
+    layer = _layer()
+    x = torch.randn(4, 32, 128, requires_grad=True)
+    mask = torch.cat([torch.ones(4, 20), torch.zeros(4, 12)], dim=1)
+
+    fused, _ = layer(x, mask, register_hook=False)
+    explicit, _ = layer(x, mask, register_hook=True)
+
+    assert torch.allclose(fused, explicit, atol=1e-5)
+
+
+def test_fused_path_does_not_retain_the_attention_map():
+    """The map was kept in memory even when no caller read it."""
+    layer = _layer()
+    x = torch.randn(2, 16, 128)
+    mask = torch.ones(2, 16)
+
+    with torch.no_grad():
+        layer(x, mask, register_hook=False)
+
+    for module in layer.modules():
+        if hasattr(module, "attn_map"):
+            assert module.attn_map is None
+
+
+def test_explicit_path_still_supplies_the_attention_map():
+    layer = _layer()
+    x = torch.randn(2, 16, 128, requires_grad=True)
+    mask = torch.ones(2, 16)
+
+    layer(x, mask, register_hook=True)
+
+    maps = [m.attn_map for m in layer.modules() if hasattr(m, "attn_map")]
+    assert any(m is not None for m in maps)
+
+
+def test_mask_fill_value_is_representable_in_half_precision():
+    """``-1e9`` is outside the fp16 range, so it cannot be the fill value.
+
+    Under fp16 autocast the previous constant overflows. The fill value must come
+    from the dtype.
+    """
+    from pyhealth.models.transformer import Attention
+
+    attention = Attention()
+    query = torch.randn(1, 1, 4, 8, dtype=torch.float16)
+    key = torch.randn(1, 1, 4, 8, dtype=torch.float16)
+    value = torch.randn(1, 1, 4, 8, dtype=torch.float16)
+    mask = torch.tensor([[[[1, 1, 0, 0]]]])
+
+    out, weights = attention(query, key, value, mask=mask)
+
+    assert torch.isfinite(out).all()
+    assert torch.isfinite(weights).all()
+    # Masked positions must receive exactly zero weight.
+    assert weights[..., 2:].abs().max().item() == 0.0
+
+
+def test_gradients_flow_through_the_fused_path():
+    layer = TransformerLayer(feature_size=64, heads=4, num_layers=2)
+    x = torch.randn(2, 16, 64)
+    out, _ = layer(x, torch.ones(2, 16))
+    out.sum().backward()
+
+    total = sum(
+        p.grad.abs().sum().item() for p in layer.parameters() if p.grad is not None
+    )
+    assert total > 0
+
+
+def test_amp_dtype_is_validated_not_silently_coerced():
+    """Any value other than "bf16" previously selected fp16 with no message.
+
+    fp16 also needs a GradScaler, so the silent path changed gradient behaviour
+    as well as precision.
+    """
+    from pyhealth.trainer import resolve_amp_dtype
+
+    assert resolve_amp_dtype("bf16") is torch.bfloat16
+    assert resolve_amp_dtype("fp16") is torch.float16
+    # The long spellings used to fall through to fp16.
+    assert resolve_amp_dtype("bfloat16") is torch.bfloat16
+    assert resolve_amp_dtype("float16") is torch.float16
+
+    for bad in ("bfloat_16", "f16", "int8", ""):
+        with pytest.raises(ValueError):
+            resolve_amp_dtype(bad)