From b4db805370832826ac0f58f7353843160dd16272 Mon Sep 17 00:00:00 2001 From: Rian354 Date: Mon, 10 Aug 2026 14:06:02 -0400 Subject: [PATCH] Use fused attention, make mask fill fp16-safe, validate the AMP dtype, expose DataLoader worker options The explicit attention path builds several (B, H, S, S) tensors for each layer and is needed only for interpretability. The mask fill value -1e9 is outside the fp16 range. Any amp_dtype other than bf16 previously selected fp16 with no message. Measurements are in the PR description. --- pyhealth/datasets/utils.py | 49 ++++++++++++-- pyhealth/models/transformer.py | 53 +++++++++++---- pyhealth/trainer.py | 39 ++++++++++- tests/test_attention_paths.py | 120 +++++++++++++++++++++++++++++++++ 4 files changed, 238 insertions(+), 23 deletions(-) create mode 100644 tests/test_attention_paths.py diff --git a/pyhealth/datasets/utils.py b/pyhealth/datasets/utils.py index ce958f627..32425fc7c 100644 --- a/pyhealth/datasets/utils.py +++ b/pyhealth/datasets/utils.py @@ -327,7 +327,13 @@ 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. @@ -335,18 +341,47 @@ def get_dataloader( 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]: diff --git a/pyhealth/models/transformer.py b/pyhealth/models/transformer.py index f678403b2..53c210f7e 100644 --- a/pyhealth/models/transformer.py +++ b/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) diff --git a/pyhealth/trainer.py b/pyhealth/trainer.py index fc264a3af..5e7c52893 100644 --- a/pyhealth/trainer.py +++ b/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() diff --git a/tests/test_attention_paths.py b/tests/test_attention_paths.py new file mode 100644 index 000000000..366fe8913 --- /dev/null +++ b/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)