Introduce FlashAttention3 support

Author: cbalioglu

src/fairseq2/datasets/_batch.py

@@ -13,6 +13,7 @@
 from typing_extensions import override
 
 from fairseq2.device import Device, SupportsDeviceTransfer
+from fairseq2.error import InvalidOperationError
 from fairseq2.nn import BatchLayout
 
 
@@ -105,7 +106,7 @@ def __init__(
             self._padding = 0
 
             for idx, seq_len in enumerate(seq_lens):
-                if seq_len < 0:
+                if seq_len < 1:
                     raise ValueError(
                         f"All lengths in `seq_lens` must be greater than or equal to 1, but the length at index {idx} is {seq_len} instead."
                     )
@@ -163,11 +164,27 @@ def as_auto_regressive(self) -> tuple[SequenceBatch, SequenceBatch]:
 
             seq_lens = self._seq_lens.copy()
 
-            seq_lens[-1] -= 1
+            if seq_lens[-1] == 1:
+                if len(seq_lens) == 1:
+                    raise InvalidOperationError(
+                        "The length of the sequence at index 0 is already 1 and cannot be trimmed to 0."
+                    )
+
+                del seq_lens[-1]
+            else:
+                seq_lens[-1] -= 1
         else:
             seqs = self._seqs[:, :-1]
 
-            seq_lens = [seq_len - 1 for seq_len in self._seq_lens]
+            seq_lens = []
+
+            for idx, seq_len in enumerate(self._seq_lens):
+                if seq_len == 1:
+                    raise InvalidOperationError(
+                        f"The length of the sequence at index {idx} is already 1 and cannot be trimmed to 0."
+                    )
+
+                seq_lens.append(seq_len - 1)
 
         batch = SequenceBatch(
             seqs, seq_lens, packed=self._packed, example=self._example
@@ -487,11 +504,27 @@ def as_auto_regressive(self) -> tuple[Seq2SeqBatch, SequenceBatch]:
 
             seq_lens = self._target_seq_lens.copy()
 
-            seq_lens[-1] -= 1
+            if seq_lens[-1] == 1:
+                if len(seq_lens) == 1:
+                    raise InvalidOperationError(
+                        "The length of the target sequence at index 0 is already 1 and cannot be trimmed to 0."
+                    )
+
+                del seq_lens[-1]
+            else:
+                seq_lens[-1] -= 1
         else:
             seqs = self._target_seqs[:, :-1]
 
-            seq_lens = [seq_len - 1 for seq_len in self._target_seq_lens]
+            seq_lens = []
+
+            for idx, seq_len in enumerate(self._target_seq_lens):
+                if seq_len == 1:
+                    raise InvalidOperationError(
+                        f"The length of the target sequence at index {idx} is already 1 and cannot be trimmed to 0."
+                    )
+
+                seq_lens.append(seq_len - 1)
 
         batch = Seq2SeqBatch(
             self._source_seqs,

src/fairseq2/datasets/instruction.py

@@ -297,7 +297,9 @@ def skip(example: dict[str, Any]) -> bool:
             "target_mask", pad_value=False
         )
 
-        collater = Collater(pad_value=0, overrides=[target_mask_collate_opts])
+        collater = Collater(
+            pad_value=tokenizer.vocab_info.pad_idx, overrides=[target_mask_collate_opts]
+        )
 
         builder.map(collater, num_parallel_calls=npc)
 

src/fairseq2/metrics/recorders/_tensorboard.py

@@ -13,6 +13,13 @@
 
 from typing_extensions import override
 
+try:
+    from torch.utils.tensorboard import SummaryWriter  # type: ignore[attr-defined]
+except ImportError:
+    _has_tensorboard = False
+else:
+    _has_tensorboard = True
+
 from fairseq2.logging import log
 from fairseq2.metrics import MetricDescriptor
 from fairseq2.registry import Provider
@@ -28,13 +35,6 @@
     NoopMetricRecorder,
 )
 
-try:
-    from torch.utils.tensorboard import SummaryWriter  # type: ignore[attr-defined]
-except ImportError:
-    has_tensorboard = False
-else:
-    has_tensorboard = True
-
 
 @final
 class TensorBoardRecorder(MetricRecorder):
@@ -51,7 +51,7 @@ def __init__(
         :param output_dir:
             The base directory under which to store the TensorBoard files.
         """
-        if not has_tensorboard:
+        if not _has_tensorboard:
             log.warning("tensorboard not found. Please install it with `pip install tensorboard`.")  # fmt: skip
 
         self._output_dir = output_dir
@@ -94,7 +94,7 @@ def record_metrics(
             ) from ex
 
     def _get_writer(self, run: str) -> SummaryWriter | None:
-        if not has_tensorboard:
+        if not _has_tensorboard:
             return None
 
         writer = self._writers.get(run)

src/fairseq2/metrics/recorders/_wandb.py

@@ -13,6 +13,13 @@
 
 from typing_extensions import override
 
+try:
+    import wandb  # type: ignore[import-not-found]
+except ImportError:
+    _has_wandb = False
+else:
+    _has_wandb = True
+
 from fairseq2.logging import log
 from fairseq2.metrics import MetricDescriptor
 from fairseq2.registry import Provider
@@ -28,13 +35,6 @@
     NoopMetricRecorder,
 )
 
-try:
-    import wandb  # type: ignore[import-not-found]
-except ImportError:
-    has_wandb = False
-else:
-    has_wandb = True
-
 
 @final
 class WandbRecorder(MetricRecorder):
@@ -57,7 +57,7 @@ def __init__(
         In order to use W&B, run `wandb login` from the command line and enter
         the API key when prompted.
         """
-        if not has_wandb:
+        if not _has_wandb:
             log.warning("wandb not found. Please install it with `pip install wandb`.")  # fmt: skip
 
             self._run = None

src/fairseq2/models/llama/_factory.py

@@ -93,8 +93,9 @@ def init_embed(embed: StandardEmbedding) -> None:
             self._init_truncated_normal(embed.weight, bias=None, std=std)
 
         return StandardEmbedding(
-            num_embeddings=config.vocab_size,
-            embedding_dim=config.model_dim,
+            config.vocab_size,
+            config.model_dim,
+            pad_idx=config.pad_idx,
             init_fn=init_embed,
         )
 

src/fairseq2/models/mistral/_factory.py

@@ -84,7 +84,7 @@ def create_embedding(self) -> Embedding:
         config = self._config
 
         return StandardEmbedding(
-            num_embeddings=config.vocab_size, embedding_dim=config.model_dim
+            config.vocab_size, config.model_dim, pad_idx=config.pad_idx
         )
 
     def create_decoder(self) -> TransformerLMDecoder:

src/fairseq2/models/s2t_transformer/_factory.py

@@ -232,8 +232,8 @@ def create_target_embedding(self) -> Embedding:
         config = self._config
 
         return StandardEmbedding(
-            num_embeddings=config.target_vocab_size,
-            embedding_dim=config.model_dim,
+            config.target_vocab_size,
+            config.model_dim,
             pad_idx=config.pad_idx,
             init_fn=init_scaled_embedding,
         )

src/fairseq2/models/transformer/__init__.py

@@ -142,6 +142,8 @@
 from fairseq2.models.transformer._sdpa._default import (
     set_default_sdpa_factory as set_default_sdpa_factory,
 )
+from fairseq2.models.transformer._sdpa._flash2 import Flash2SDPA as Flash2SDPA
+from fairseq2.models.transformer._sdpa._flash3 import Flash3SDPA as Flash3SDPA
 from fairseq2.models.transformer._sdpa._naive import NaiveSDPA as NaiveSDPA
 from fairseq2.models.transformer._sdpa._naive import (
     naive_scaled_dot_product_attention as naive_scaled_dot_product_attention,

src/fairseq2/models/transformer/_factory.py

@@ -96,8 +96,8 @@ def create_embedding(self) -> Embedding:
         config = self._config
 
         return StandardEmbedding(
-            num_embeddings=config.vocab_size,
-            embedding_dim=config.model_dim,
+            config.vocab_size,
+            config.model_dim,
             pad_idx=config.pad_idx,
             init_fn=init_scaled_embedding,
         )