diff --git a/pyproject.toml b/pyproject.toml index 29c56c4..b7a11af 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "mblm" -version = "0.4.0" +version = "0.4.1" description = "Multiscale Byte Language Model" authors = [ { name = "Eric Egli", email = "eric.christian.egli@ibm.com" }, diff --git a/src/mblm/model/embeddings.py b/src/mblm/model/embeddings.py index 19a8453..aea0d6e 100644 --- a/src/mblm/model/embeddings.py +++ b/src/mblm/model/embeddings.py @@ -22,8 +22,8 @@ MBLM_TOKEN_EMB_MIGRATION: set[str] = { - "token_embs_rev.0.weight", - "token_embs_rev.1.0.weight", + "token_embs_rev.0.embedding.weight", + "token_embs_rev.1.embedding.weight", "to_logits.weight", "to_logits.bias", } diff --git a/src/mblm/model/mblm.py b/src/mblm/model/mblm.py index d8fe611..6d9480a 100644 --- a/src/mblm/model/mblm.py +++ b/src/mblm/model/mblm.py @@ -34,6 +34,7 @@ from mblm.model.block import StageBlock from mblm.model.config import MBLMEncoderModelConfig, MBLMModelConfig, MBLMReturnType +from mblm.model.multi_stage_token_embedding import MultiStageTokenEmbedding from mblm.model.utils import RoPE, gumbel_sample, top_k from mblm.utils.stream import ByteStreamer @@ -88,8 +89,11 @@ def __init__(self, cfg: MBLMModelConfig): cfg.hidden_dims, cfg.seq_lens, stage_blocks ) - self.token_embs_rev = self._init_token_embeddings( - cfg.hidden_dims, cfg.seq_lens, cfg.num_tokens, cfg.pad_token_id + self.token_embs_rev = MultiStageTokenEmbedding.build( + model_dims=cfg.hidden_dims, + seq_lens=cfg.seq_lens, + vocab_size=cfg.num_tokens, + pad_token_id=cfg.pad_token_id, ) self.stage_models, self.to_next_stage_proj = self._init_models_at_stages( @@ -98,6 +102,8 @@ def __init__(self, cfg: MBLMModelConfig): self.to_logits = nn.Linear(cfg.hidden_dims[-1], cfg.num_tokens) + self._warned_single_tensor_embeds: bool = False + @classmethod def _init_positional_embeddings( cls, @@ -126,41 +132,6 @@ def _init_positional_embeddings( ) return modules - @classmethod - def _init_token_embeddings( - cls, - model_dims: Sequence[int], - seq_lens: Sequence[int], - vocab_size: int, - pad_token_id: int, - ) -> nn.ModuleList: - """ - Embed the tokens for each stage (in reverse order). - """ - local_dim = model_dims[-1] - token_embs_rev = nn.ModuleList( - [nn.Embedding(vocab_size, local_dim, padding_idx=pad_token_id)] - ) - patch_size = 1 - for model_dim, seq_len in zip( - # all except the local model - reversed(model_dims[:-1]), # (D_n-1, ..., D_1) - reversed(seq_lens[1:]), # (P_2, ..., P_n) - ): - # for the global models, fuse the embedding and patch projection - # step - patch_size *= seq_len - token_embs_rev.append( - nn.Sequential( - nn.Embedding(vocab_size, local_dim, padding_idx=pad_token_id), - Rearrange("... r d -> ... (r d)"), - nn.LayerNorm(patch_size * local_dim), - nn.Linear(patch_size * local_dim, model_dim), - nn.LayerNorm(model_dim), - ) - ) - return token_embs_rev - @classmethod def _init_models_at_stages( cls, @@ -221,16 +192,16 @@ def forward_empty(self, batch_size: int) -> torch.Tensor: @overload def forward( self, - input_ids: torch.Tensor, - *, + input_ids: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor | Sequence[torch.Tensor]] = None, return_type: Literal[MBLMReturnType.LOSS_LOGITS] = ..., loss_mask: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: ... @overload def forward( self, - input_ids: torch.Tensor, - *, + input_ids: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor | Sequence[torch.Tensor]] = None, return_type: Literal[ MBLMReturnType.LOSS, MBLMReturnType.LOGITS, MBLMReturnType.HIDDEN_STATE ] = ..., @@ -239,8 +210,8 @@ def forward( def forward( self, - input_ids: torch.Tensor, - *, + input_ids: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor | Sequence[torch.Tensor]] = None, return_type: MBLMReturnType = MBLMReturnType.LOSS, loss_mask: torch.Tensor | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: @@ -248,7 +219,17 @@ def forward( A single forward pass. Args: - input_ids: The token ids as torch.LongTensor in shape (B, L) + input_ids: The token ids as torch.LongTensor in shape (B, L). + inputs_embeds: Precomputed embeddings as float tensor or sequence of tensors. + Must be mutually exclusive with input_ids. Can be either: + - Single tensor: Will be reused across all stages (bypasses embedding lookups). + The tensor must be pre-organized into the hierarchical structure: + (B, P1', P2, ..., Pn, Dn) for multi-stage or (B, L, D) for single-stage. + NOTE: This will produce different results than input_ids because each stage + has its own embedding table. + - Sequence of tensors: One tensor per stage (in reverse order: local to global). + This allows exact replication of the input_ids path. + Expected shapes: [local_embs (B, P1', P2, ..., Pn, Dn), ..., global_embs (B, P1', D1)] return_type: What to return - the loss, the logits or both. loss_mask: An optional masking tensor that enables interpolation between self-supervised and supervised learning. It determines which tokens in the @@ -262,54 +243,139 @@ def forward( By default, not providing a `loss_mask` is equivalent to a `loss_mask` consisting of all 1 """ - batch_size = input_ids.shape[0] + if (input_ids is None) == (inputs_embeds is None): + raise ValueError("Pass exactly one of input_ids or inputs_embeds.") + + embeds_list: Optional[Sequence[torch.Tensor]] = None + + if input_ids is not None: + device = input_ids.device + batch_size = input_ids.shape[0] + + assert input_ids.ndim in {2, self.num_stages + 1} + + if input_ids.numel() == 0: + return self.forward_empty(batch_size) + + if loss_mask is not None: + assert loss_mask.shape == input_ids.shape + + flattened_dims = input_ids.ndim == 2 + flat_seq_len = input_ids.shape[-1] + + # if the input is given as (B, L), reshape and distribute it among the + # hierarchy sequence lengths, filling up the inner dimensions + # first. padding is applied so that the output shape is (B, P_1', P_2, + # ..., P_n). for the largest possible L == prod(seq_lens), P_1' = P_1. + # in all other cases, P_1' < P_1, meaning no padding is applied to the + # global sequence P_1. + # + # here's two examples for a model model (P_1, P_2, P_3) = (5, 4, 3): + # + # input: (B, L) = (1, 13), output: (B, P_1', P_2, P_3) = (1, 2, 4, 3) + # input: (B, L) = (1, 60), output: (B, P_1, P_2, P_3) = (1, 5, 4, 3) + # + # in the 2nd example, L = prod(seq_lens) = 5 * 4 * 3 = 60 = P_1 = P_1' + if flattened_dims: + # pad/fill up all inner sequence lengths (all except most global) + local_seq_lens = self.seq_lens[1:] + multiple_of = math.prod(local_seq_lens) + # use the complement of modulo - the difference to the next multiple + # of multiple_of - to infer the right padding length + padding = -flat_seq_len % multiple_of + input_ids = F.pad(input_ids, (0, padding), value=self.pad_token_id) + # reshape and infer the P_1' dimension + input_ids = input_ids.reshape(batch_size, -1, *local_seq_lens) + + # make sure the above condition holds, i.e., P_1' <= P_1 + _P_1_prime, _P_1 = input_ids.shape[1], self.seq_lens[0] # noqa: N806 + fixed_global_patch_encoding = isinstance(self.pos_embs[0], nn.Embedding) + if fixed_global_patch_encoding: + assert _P_1_prime <= _P_1, ( + f"Because you are using a fixed global patch embedding, " + f"the input sequence length ({_P_1_prime}) " + f"must be less than the first tuple element of seq_lens ({_P_1})" + ) + elif inputs_embeds is not None: + flattened_dims = False + if return_type != MBLMReturnType.HIDDEN_STATE: + raise ValueError( + f"Forward pass with return type {return_type} has not been tested and is not supported " + f"when providing `inputs_embeds`." + ) - assert input_ids.ndim in {2, self.num_stages + 1} + # Check if inputs_embeds is a sequence of tensors or a single tensor + is_sequence = isinstance(inputs_embeds, (list, tuple)) - if input_ids.numel() == 0: - return self.forward_empty(batch_size) + if is_sequence: + # Sequence of tensors: one per stage (in reverse order: local to global) + # Type narrowing: we know inputs_embeds is a Sequence here + inputs_embeds_seq: Sequence[torch.Tensor] = cast( + Sequence[torch.Tensor], inputs_embeds + ) + assert len(inputs_embeds_seq) == self.num_stages, ( + f"`inputs_embeds` sequence must have {self.num_stages} tensors (one per stage), " + f"got {len(inputs_embeds_seq)}" + ) + # Validate each tensor in the sequence + device = inputs_embeds_seq[0].device + for stage_idx, stage_embeds in enumerate(inputs_embeds_seq): + assert isinstance( + stage_embeds, torch.Tensor + ), f"`inputs_embeds[{stage_idx}]` must be a torch.Tensor" + assert ( + stage_embeds.device == device + ), "All tensors in `inputs_embeds` must be on the same device" + + embeds_list = inputs_embeds_seq + else: + # Single tensor: will be reused across all stages + # Type narrowing: we know inputs_embeds is a Tensor here + inputs_embeds_single = cast(torch.Tensor, inputs_embeds) + device = inputs_embeds_single.device + + # Warn once per instance that single tensor cannot reproduce input_ids results + if not self._warned_single_tensor_embeds and self.num_stages > 1: + import warnings + + warnings.warn( + "Passing a single tensor to `inputs_embeds` bypasses per-stage embedding tables " + "and cannot fully reproduce the results of passing `input_ids`. " + "For exact consistency, pass a sequence of tensors (one per stage) instead.", + UserWarning, + stacklevel=2, + ) + self._warned_single_tensor_embeds = True - if loss_mask is not None: - assert loss_mask.shape == input_ids.shape + assert inputs_embeds_single.ndim == self.num_stages + 2, ( + f"`inputs_embeds` must be nested with {self.num_stages} hierarchy dims " + f"(B + P1'..Pn + D), got shape {tuple(inputs_embeds_single.shape)}" + ) + final_hidden_dim = self.start_tokens[-1].numel() + assert inputs_embeds_single.shape[-1] == final_hidden_dim, ( + f"Last dim of `inputs_embeds` must equal final hidden dim ({final_hidden_dim}), " + f"got {inputs_embeds_single.shape[-1]}" + ) + # Inner hierarchy dims must match config exactly + for k, expected in enumerate(self.seq_lens[1:], start=2): + got = inputs_embeds_single.shape[k] + assert ( + got == expected + ), f"`inputs_embeds` inner dim at stage {k - 1} must be {expected}, got {got}" + # Global fixed pos-emb requires P1' <= P1 + if isinstance(self.pos_embs[0], nn.Embedding): + assert inputs_embeds_single.shape[1] <= self.seq_lens[0], ( + f"With fixed global positional embedding, P1'={inputs_embeds_single.shape[1]} must " + f"be <= P1={self.seq_lens[0]}" + ) + # `flat_seq_len` is irrelevant in this branch (we do not compute loss/logits) + flat_seq_len = 0 # keep a defined name for readability - flattened_dims = input_ids.ndim == 2 - flat_seq_len = input_ids.shape[-1] + token_embs_at_stages = [torch.empty(0) for _ in range(self.num_stages)] - # if the input is given as (B, L), reshape and distribute it among the - # hierarchy sequence lengths, filling up the inner dimensions - # first. padding is applied so that the output shape is (B, P_1', P_2, - # ..., P_n). for the largest possible L == prod(seq_lens), P_1' = P_1. - # in all other cases, P_1' < P_1, meaning no padding is applied to the - # global sequence P_1. - # - # here's two examples for a model model (P_1, P_2, P_3) = (5, 4, 3): - # - # input: (B, L) = (1, 13), output: (B, P_1', P_2, P_3) = (1, 2, 4, 3) - # input: (B, L) = (1, 60), output: (B, P_1, P_2, P_3) = (1, 5, 4, 3) - # - # in the 2nd example, L = prod(seq_lens) = 5 * 4 * 3 = 60 = P_1 = P_1' - if flattened_dims: - # pad/fill up all inner sequence lengths (all except most global) - local_seq_lens = self.seq_lens[1:] - multiple_of = math.prod(local_seq_lens) - # use the complement of modulo - the difference to the next multiple - # of multiple_of - to infer the right padding length - padding = -flat_seq_len % multiple_of - input_ids = F.pad(input_ids, (0, padding), value=self.pad_token_id) - # reshape and infer the P_1' dimension - input_ids = input_ids.reshape(batch_size, -1, *local_seq_lens) - - # make sure the above condition holds, i.e., P_1' <= P_1 - _P_1_prime, _P_1 = input_ids.shape[1], self.seq_lens[0] # noqa: N806 - fixed_global_patch_encoding = isinstance(self.pos_embs[0], nn.Embedding) - if fixed_global_patch_encoding: - assert _P_1_prime <= _P_1, ( - f"Because you are using a fixed global patch embedding, " - f"the input sequence length ({_P_1_prime}) " - f"must be less than the first tuple element of seq_lens ({_P_1})" - ) + ids_buf = input_ids + embeds_buf = inputs_embeds if inputs_embeds is not None and not is_sequence else None - token_embs_at_stages = [torch.empty(0) for _ in range(self.num_stages)] # at this stage, we're working with nested ids - hence, embed the bytes # for each stage in reverse order, starting from the local and ending at # the most global model. at each stage, add positional embeddings and @@ -324,12 +390,17 @@ def forward( reversed(self.pos_embs), self.token_embs_rev, ): - stage_token_embs: torch.Tensor = token_emb(input_ids) + # If we have a sequence of embeddings, use them directly (they're already embedded+projected) + if embeds_list is not None: + # embeds_list is in reverse order (local to global), matching stage_idx iteration + stage_token_embs = embeds_list[self.num_stages - 1 - stage_idx] + else: + # Normal path: embed the ids or embeds_buf + stage_token_embs = token_emb(ids_buf, embeds_buf) stage_seq_len = stage_token_embs.shape[-2] - if isinstance(pos_emb, nn.Embedding): positions: torch.Tensor = pos_emb( - torch.arange(stage_seq_len, device=input_ids.device), + torch.arange(stage_seq_len, device=device), ) stage_token_embs = stage_token_embs + positions elif isinstance(pos_emb, RoPE): @@ -346,7 +417,12 @@ def forward( # skip rearranging for the most local model if stage_idx == self.num_stages - 1: continue - input_ids = rearrange(input_ids, "... m n -> ... (m n)") + + if ids_buf is not None: + ids_buf = rearrange(ids_buf, "... m n -> ... (m n)") + elif embeds_list is None: + # Only rearrange embeds_buf if we're using a single tensor (not a sequence) + embeds_buf = rearrange(embeds_buf, "... m n d -> ... (m n) d") # type: ignore # initials prev_stage_tokens_repr: torch.Tensor | None = None @@ -475,8 +551,8 @@ def forward( # (ensured by the datasets/dataloaders) as well as patch-padding # (ensured by bootstrapping MBLM with the right pad token id) loss_tensor: torch.Tensor = F.cross_entropy( - preds, # (B, V, L) - targets, # (B, L) + preds, + targets, # type: ignore ignore_index=self.pad_token_id, reduction="none", ) @@ -574,15 +650,20 @@ def __init__(self, config: MBLMEncoderModelConfig, **kwargs): def forward( self, - masked_input_ids: torch.Tensor, + masked_input_ids: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, mask: Optional[torch.Tensor] = None, labels: Optional[torch.Tensor] = None, return_type: MBLMReturnType = MBLMReturnType.HIDDEN_STATE, ): if return_type == MBLMReturnType.HIDDEN_STATE: - return self.mblm.forward(masked_input_ids, return_type=MBLMReturnType.HIDDEN_STATE) + return self.mblm.forward( + masked_input_ids, inputs_embeds, return_type=MBLMReturnType.HIDDEN_STATE + ) - logits = self.mblm.forward(masked_input_ids, return_type=MBLMReturnType.LOGITS) + logits = self.mblm.forward( + masked_input_ids, inputs_embeds, return_type=MBLMReturnType.LOGITS + ) if return_type == MBLMReturnType.LOGITS: return logits # ignore non mask token in the loss computation, this is used with the ignore_index parameter of cross_entropy diff --git a/src/mblm/model/multi_stage_token_embedding.py b/src/mblm/model/multi_stage_token_embedding.py new file mode 100644 index 0000000..c4b5c72 --- /dev/null +++ b/src/mblm/model/multi_stage_token_embedding.py @@ -0,0 +1,207 @@ +from __future__ import annotations + +__copyright__ = """MIT License + +Copyright (c) 2024 - IBM Research + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE.""" + +from typing import Optional, Sequence + +import torch +import torch.nn as nn +from einops.layers.torch import Rearrange + + +class _StageTokenEmbedding(nn.Module): + """ + Per-stage embedding module that supports either input_ids or inputs_embeds. + + For local stage (is_local=True): + - If input_ids is given: returns embedding(input_ids) -> [..., L, local_dim] + - If inputs_embeds is given: returns inputs_embeds as-is (expects [..., L, local_dim]) + + For global stage (is_local=False): + - If input_ids is given: embedding -> [..., R, local_dim] -> flatten/proj -> [..., model_dim] + - If inputs_embeds is given: expects [..., R, local_dim] -> flatten/proj -> [..., model_dim] + + In all cases, the projection (flatten + LN + Linear + LN) reuses the same weights, + so switching between input_ids and inputs_embeds produces identical results + (provided inputs_embeds come from the same embedding lookup). + """ + + def __init__( + self, + *, + vocab_size: int, + pad_token_id: int, + local_dim: int, + is_local: bool, + patch_size: int = 1, # product of downstream seq_lens + model_dim: Optional[int] = None, # required if is_local=False + ) -> None: + super().__init__() + self.is_local = is_local + self.local_dim = int(local_dim) + self.patch_size = int(patch_size) + self.model_dim = None if model_dim is None else int(model_dim) + + # The (shared) token embedding for input_ids path + self.embedding = nn.Embedding( + num_embeddings=int(vocab_size), + embedding_dim=self.local_dim, + padding_idx=pad_token_id, + ) + + if self.is_local: + # No projection for local stage + self._post = nn.Identity() + else: + if self.model_dim is None: + raise ValueError("model_dim must be provided for global stages.") + flat_dim = self.patch_size * self.local_dim + self._post = nn.Sequential( # type: ignore + Rearrange("... r d -> ... (r d)"), + nn.LayerNorm(flat_dim), + nn.Linear(flat_dim, self.model_dim), + nn.LayerNorm(self.model_dim), + ) + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """ + Exactly one of (input_ids, inputs_embeds) must be provided. + + Shapes: + - local stage: + input_ids: [..., L] -> [..., L, local_dim] + inputs_embeds: [..., L, local_dim] (returned unchanged) + - global stage: + input_ids: [..., R] -> [..., R, local_dim] -> proj -> [..., model_dim] + inputs_embeds: [..., R, local_dim] -> proj -> [..., model_dim] + """ + if (input_ids is None) == (inputs_embeds is None): + raise ValueError("Pass exactly one of input_ids or inputs_embeds.") + + if input_ids is not None: + # IDs -> embedding lookup + x = self.embedding(input_ids) # local: [..., L, d]; global: [..., R, d] + else: + # bypass embedding layer; reuse projection weights + x = inputs_embeds + # --- Shape validation for inputs_embeds --- + if not isinstance(x, torch.Tensor): + raise RuntimeError("inputs_embeds must be a torch.Tensor.") + if x.ndim < 2: + raise RuntimeError(f"inputs_embeds must be at least 2D (got {x.ndim}D).") + # Local stage expects [..., L, local_dim] + if self.is_local: + if x.shape[-1] != self.local_dim: + raise RuntimeError( + f"Local stage expects inputs_embeds[..., local_dim={self.local_dim}], " + f"but got last dim={x.shape[-1]}." + ) + else: + # Global stage expects [..., R, local_dim] with R == self.patch_size + if x.ndim < 2: + raise RuntimeError( + "Global stage expects inputs_embeds with at least 2 dims [..., R, d]." + ) + if x.shape[-2] != self.patch_size: + raise RuntimeError( + f"Global stage expects inputs_embeds[..., R={self.patch_size}, d], " + f"but got R={x.shape[-2]}." + ) + if x.shape[-1] != self.local_dim: + raise RuntimeError( + f"Global stage expects last dim local_dim={self.local_dim}, " + f"but got {x.shape[-1]}." + ) + + if self.is_local: + # local stage has no projection + return x + else: + # global stage: flatten(r, d) and project + return self._post(x) + + +class MultiStageTokenEmbedding: + """ + Factory that replicates the behavior of _init_token_embeddings but returns + per-stage modules that accept both input_ids and inputs_embeds. + + The returned ModuleList is in REVERSE order. + """ + + @classmethod + def build( + cls, + *, + model_dims: Sequence[int], + seq_lens: Sequence[int], + vocab_size: int, + pad_token_id: int, + ) -> nn.ModuleList: + """ + Replicates the original initializer semantics (reverse order, fused global projection). + + Returns: + nn.ModuleList in reverse stage order: + [ local_stage_module, global_n-1, ..., global_1 ] + """ + model_dims = list(model_dims) + seq_lens = list(seq_lens) + + local_dim = model_dims[-1] # local model hidden dim + token_embs_rev = nn.ModuleList( + [ + _StageTokenEmbedding( + vocab_size=vocab_size, + pad_token_id=pad_token_id, + local_dim=local_dim, + is_local=True, # last stage is local + patch_size=1, + model_dim=None, + ) + ] + ) + + patch_size = 1 + # iterate over global models in reverse, accumulating patch_size + for model_dim, seq_len in zip( + reversed(model_dims[:-1]), # (D_{n-1}, ..., D_1) + reversed(seq_lens[1:]), # (P_2, ..., P_n) + ): + patch_size *= int(seq_len) + token_embs_rev.append( + _StageTokenEmbedding( + vocab_size=vocab_size, + pad_token_id=pad_token_id, + local_dim=local_dim, + is_local=False, + patch_size=patch_size, + model_dim=int(model_dim), + ) + ) + + return token_embs_rev diff --git a/tests/unit/model/test_mblm.py b/tests/unit/model/test_mblm.py index af7e667..ea7a630 100644 --- a/tests/unit/model/test_mblm.py +++ b/tests/unit/model/test_mblm.py @@ -161,7 +161,9 @@ def test_masked_mblm_fully_masked_is_nan( mask = torch.zeros_like(input_ids) mask = mask.to(torch.bool) masked_input[mask] = self.mask_token_id - loss = masked_model.forward(masked_input, mask, input_ids, return_type=MBLMReturnType.LOSS) + loss = masked_model.forward( + masked_input, mask=mask, labels=input_ids, return_type=MBLMReturnType.LOSS + ) assert loss.isnan().item(), f"Got {loss.isnan().item()}" def test_masked_mblm_partially_masked_is_float( @@ -176,7 +178,9 @@ def test_masked_mblm_partially_masked_is_float( mask = torch.rand_like(input_ids, dtype=torch.float) < 0.15 mask = mask.to(torch.bool) masked_input[mask] = self.mask_token_id - loss = masked_model.forward(masked_input, mask, input_ids, return_type=MBLMReturnType.LOSS) + loss = masked_model.forward( + masked_input, mask=mask, labels=input_ids, return_type=MBLMReturnType.LOSS + ) assert loss.dtype == torch.float and loss.item() > 0.0 @pytest.mark.parametrize("batch", [1, 3]) @@ -191,10 +195,10 @@ def test_masked_mblm_return_type_shape(self, batch): mask = mask.to(torch.bool) masked_input[mask] = self.mask_token_id loss, logit = masked_model.forward( - masked_input, mask, input_ids, return_type=MBLMReturnType.LOSS_LOGITS + masked_input, mask=mask, labels=input_ids, return_type=MBLMReturnType.LOSS_LOGITS ) hidden_state = masked_model.forward( - masked_input, mask, input_ids, return_type=MBLMReturnType.HIDDEN_STATE + masked_input, mask=mask, labels=input_ids, return_type=MBLMReturnType.HIDDEN_STATE ) assert loss.size() == torch.Size([]) assert logit.size() == torch.Size([batch, input_len, self.mblm_conf.num_tokens]) @@ -212,13 +216,13 @@ def test_masked_mblm_combined_return(self, batch): mask = mask.to(torch.bool) masked_input[mask] = self.mask_token_id loss, logits = masked_model.forward( - masked_input, mask, input_ids, return_type=MBLMReturnType.LOSS_LOGITS + masked_input, mask=mask, labels=input_ids, return_type=MBLMReturnType.LOSS_LOGITS ) loss_only = masked_model.forward( - masked_input, mask, input_ids, return_type=MBLMReturnType.LOSS + masked_input, mask=mask, labels=input_ids, return_type=MBLMReturnType.LOSS ) logits_only = masked_model.forward( - masked_input, mask, input_ids, return_type=MBLMReturnType.LOGITS + masked_input, mask=mask, labels=input_ids, return_type=MBLMReturnType.LOGITS ) assert logits.shape == logits_only.shape assert loss.shape == loss_only.shape @@ -247,17 +251,23 @@ def test_masked_mblm_input_seq_len(self, batch): if current_input_len > max_input_length: with pytest.raises(AssertionError): masked_model.forward( - masked_input, mask, input_ids, return_type=MBLMReturnType.LOSS_LOGITS + masked_input, + mask=mask, + labels=input_ids, + return_type=MBLMReturnType.LOSS_LOGITS, ) else: loss, logits = masked_model.forward( - masked_input, mask, input_ids, return_type=MBLMReturnType.LOSS_LOGITS + masked_input, + mask=mask, + labels=input_ids, + return_type=MBLMReturnType.LOSS_LOGITS, ) loss_only = masked_model.forward( - masked_input, mask, input_ids, return_type=MBLMReturnType.LOSS + masked_input, mask=mask, labels=input_ids, return_type=MBLMReturnType.LOSS ) logits_only = masked_model.forward( - masked_input, mask, input_ids, return_type=MBLMReturnType.LOGITS + masked_input, mask=mask, labels=input_ids, return_type=MBLMReturnType.LOGITS ) assert logits.shape == logits_only.shape assert loss.shape == loss_only.shape diff --git a/tests/unit/model/test_mblm_inputs_embeds.py b/tests/unit/model/test_mblm_inputs_embeds.py new file mode 100644 index 0000000..1d65335 --- /dev/null +++ b/tests/unit/model/test_mblm_inputs_embeds.py @@ -0,0 +1,314 @@ +import math + +import pytest +import torch +import torch.nn.functional +from einops import rearrange + +from mblm import MBLM, MBLMModelConfig, MBLMReturnType, TransformerBlock +from mblm.model.config import MBLMEncoderModelConfig +from mblm.model.mblm import MBLMEncoder + + +def pad_and_reshape_inputs_embeds( + flat_embs: torch.Tensor, seq_lens: tuple[int, ...] +) -> tuple[torch.Tensor, int, int]: + """ + Convert (batch_size, seq_len, dim_n) -> (batch_size, p1_prime, p2, ..., pn, dim_n) + by padding to the nearest multiple of prod(seq_lens[1:]) and reshaping. + + Returns: (nested_embs, p1_prime, padding_len) + + For single-stage (len(seq_lens) == 1), this just returns (batch_size, p1_prime=seq_len, dim_n). + """ + assert flat_embs.ndim == 3, "Expected flat inputs_embeds of shape (B, L, Dn)" + batch_size, seq_len, dim_n = flat_embs.shape + inner = seq_lens[1:] + multiple_of = math.prod(inner) if len(inner) > 0 else 1 + padding_len = (-seq_len) % multiple_of + + if padding_len: + # pad along sequence dim (dim=1) on the right with zeros in emb space + flat_embs = torch.nn.functional.pad(flat_embs, (0, 0, 0, padding_len)) + + seq_len_padded = flat_embs.shape[1] + p1_prime = seq_len_padded // max(multiple_of, 1) + + if len(seq_lens) == 1: + nested = flat_embs.view(batch_size, p1_prime, dim_n) + else: + nested = flat_embs.view(batch_size, p1_prime, *inner, dim_n) + + return nested, p1_prime, padding_len + + +class TestMBLMInputsEmbeds: + """Test class.""" + + num_tokens = 256 + 1 + pad_token_id = 256 + num_attn_heads = 8 + dim_attn_heads = 64 + ff_mult = 4 + dropout = 0 + use_rot_emb = True + use_flash_attn = False + + def _build_model(self, model_dims: tuple[int, ...], seq_lens: tuple[int, ...]) -> MBLM: + return MBLM( + MBLMModelConfig( + num_tokens=self.num_tokens, + hidden_dims=model_dims, + seq_lens=seq_lens, + pad_token_id=self.pad_token_id, + num_layers=(1,) * len(model_dims), + train_checkpoint_chunks=None, + block=[ + TransformerBlock( + attn_head_dims=self.dim_attn_heads, + attn_num_heads=self.num_attn_heads, + attn_dropout=self.dropout, + ff_multiplier=self.ff_mult, + ff_dropout=self.dropout, + pos_emb_type="fixed", + attn_use_rot_embs=self.use_rot_emb, + use_flash_attn=self.use_flash_attn, + ) + ] + * len(model_dims), + ) + ) + + def test_inputs_embeds_only_allows_hidden_state_single_stage(self): + """ + Inputs_embeds is supported only with return_type == HIDDEN_STATE. + Other return types must raise. + """ + model_dims = (64,) + seq_lens = (9,) + mblm = self._build_model(model_dims, seq_lens) + mblm.eval() + + batch_size, seq_len, dim_n = 2, 7, model_dims[-1] + flat_embs = torch.randn(batch_size, seq_len, dim_n) + + # For single-stage, flat (B, L, Dn) is already (B, p1_prime, Dn) with p1_prime = L + nested = flat_embs + + with torch.no_grad(): + _ = mblm.forward( + input_ids=None, + inputs_embeds=nested, + return_type=MBLMReturnType.HIDDEN_STATE, + ) + + with pytest.raises(ValueError): + _ = mblm.forward(inputs_embeds=nested, return_type=MBLMReturnType.LOGITS) + + with pytest.raises(ValueError): + _ = mblm.forward(inputs_embeds=nested, return_type=MBLMReturnType.LOSS) + + def test_inputs_embeds_hidden_state_shape_single_stage(self): + """ + For single-stage, output should be (batch_size, 1 + p1_prime, dim_n) where p1_prime = seq_len. + """ + model_dims = (96,) + seq_lens = (11,) + mblm = self._build_model(model_dims, seq_lens) + mblm.eval() + + batch_size, seq_len, dim_n = 3, 8, model_dims[-1] + flat_embs = torch.randn(batch_size, seq_len, dim_n) + # single-stage: nested == flat along (B, p1_prime, Dn) + nested = flat_embs # p1_prime = seq_len + + with torch.no_grad(): + out = mblm.forward(inputs_embeds=nested, return_type=MBLMReturnType.HIDDEN_STATE) + + assert out.shape == torch.Size( + [batch_size, 1 + seq_len, dim_n] + ), f"got {out.shape}, expected {(batch_size, 1 + seq_len, dim_n)}" + + def test_encoder_hidden_state_matches_mblm_single_stage(self): + """ + Encoder should delegate to the same network. Synchronize weights to + assert numerical equality. + """ + model_dims = (128,) + seq_lens = (13,) + cfg = MBLMModelConfig( + num_tokens=self.num_tokens, + hidden_dims=model_dims, + seq_lens=seq_lens, + pad_token_id=self.pad_token_id, + num_layers=(1,), + train_checkpoint_chunks=None, + block=[ + TransformerBlock( + attn_head_dims=self.dim_attn_heads, + attn_num_heads=self.num_attn_heads, + attn_dropout=self.dropout, + ff_multiplier=self.ff_mult, + ff_dropout=self.dropout, + pos_emb_type="fixed", + attn_use_rot_embs=self.use_rot_emb, + use_flash_attn=self.use_flash_attn, + ) + ], + ) + + mblm = MBLM(cfg) + encoder = MBLMEncoder( + MBLMEncoderModelConfig(mask_token_id=self.pad_token_id + 1, mblm_config=cfg) + ) + encoder.mblm.load_state_dict(mblm.state_dict()) # sync weights + + mblm.eval() + encoder.eval() + + batch_size, seq_len, dim_n = 2, 5, model_dims[-1] + flat_embs = torch.randn(batch_size, seq_len, dim_n) + nested = flat_embs # single-stage + + with torch.no_grad(): + h_mblm = mblm.forward(inputs_embeds=nested, return_type=MBLMReturnType.HIDDEN_STATE) + h_encoder = encoder.forward( + inputs_embeds=nested, return_type=MBLMReturnType.HIDDEN_STATE + ) + + assert h_mblm.shape == h_encoder.shape + assert torch.allclose( + h_mblm, h_encoder, atol=1e-5 + ), "Encoder hidden states differ from MBLM after syncing weights" + + def test_inputs_embeds_nested_multistage(self): + """ + Test inputs_embeds with multistage model + """ + model_dims = (128, 64) + seq_lens = (5, 4) # prod(inner)=4 + mblm = self._build_model(model_dims, seq_lens) + mblm.eval() + + batch_size, seq_len, dim_n = ( + 1, + 7, + model_dims[-1], + ) # choose L <= P1 * prod(inner) to keep p1' <= P1 + flat_embs = torch.randn(batch_size, seq_len, dim_n) + + nested, _, _ = pad_and_reshape_inputs_embeds(flat_embs, seq_lens) + + with torch.no_grad(): + output = mblm.forward(inputs_embeds=nested, return_type=MBLMReturnType.HIDDEN_STATE) + + assert output.shape == ( + nested.shape[0], + nested.shape[1], + nested.shape[2] + 1, + model_dims[-1], + ) + + @pytest.mark.parametrize( + "model_dims, seq_lens", + [((64,), (9,)), ((128, 64), (5, 4))], + ) + def test_ids_vs_inputs_embeds_consistency_end2end(self, model_dims, seq_lens): + """ + End-to-end consistency: + 1) Forward with nested input_ids -> h_ids + 2) Compute local (most local stage) inputs_embeds EXACTLY as the model does: + local_inputs_embeds = mblm.token_embs_rev[0](nested_ids, None) + (i.e., call the local embedding on the NESTED ids, not flat) + -> forward(..., inputs_embeds) -> h_embs + 3) Normalize (drop final-stage start token, flatten, slice to original seq_len) and compare. + """ + mblm = self._build_model(model_dims, seq_lens) + mblm.eval() + + batch_size = 2 + p1 = seq_lens[0] + inner = seq_lens[1:] + prod_inner = math.prod(inner) if inner else 1 + + # Choose a valid global patch count so that p1_prime <= p1 (fixed positional embeddings at global) + p1_prime = min(2, p1) + seq_len = p1_prime * prod_inner # original token count without any start tokens + + # 1) Build nested input_ids and run ids path + nested_shape = (batch_size, p1_prime, *inner) if inner else (batch_size, p1_prime) + input_ids_nested = torch.randint(0, self.num_tokens, size=nested_shape, dtype=torch.long) + + with torch.no_grad(): + h_ids = mblm.forward( + input_ids=input_ids_nested, + inputs_embeds=None, + return_type=MBLMReturnType.HIDDEN_STATE, + ) + + # 2) Compute inputs_embeds for ALL stages exactly as forward would: + with torch.no_grad(): + inputs_embeds_list = [] + ids_buf = input_ids_nested + embeds_buf = None + + # Process stages in reverse order (local to global), matching the forward loop + for stage_idx in range(len(seq_lens) - 1, -1, -1): + token_emb = mblm.token_embs_rev[len(seq_lens) - 1 - stage_idx] + # Compute embeddings for this stage + stage_embeds = token_emb(ids_buf, embeds_buf) + inputs_embeds_list.append(stage_embeds) + + # Rearrange for next (more global) stage, except for the most local + if stage_idx < len(seq_lens) - 1: + if ids_buf is not None: + ids_buf = rearrange(ids_buf, "... m n -> ... (m n)") + else: + embeds_buf = rearrange(embeds_buf, "... m n d -> ... (m n) d") + + # forward on inputs_embeds sequence + h_embs = mblm.forward( + input_ids=None, + inputs_embeds=inputs_embeds_list, + return_type=MBLMReturnType.HIDDEN_STATE, + ) + + # 3) Normalize both to (B, L, Dn): drop final-stage start token if present, + # flatten hierarchical dims, and slice to original seq_len + def normalize_hidden(hidden: torch.Tensor) -> torch.Tensor: + if len(seq_lens) == 1: + # Single-stage: shapes are (B, S, D) + # If start token included: S == p1_prime + 1 -> drop the first token on that axis + if hidden.shape[1] == p1_prime + 1: + hidden = hidden[:, 1:, :] + elif hidden.shape[1] != p1_prime: + raise AssertionError( + f"Unexpected single-stage shape {tuple(hidden.shape)} with p1'={p1_prime}" + ) + # Already (B, p1_prime, D) + out_hidden = hidden + else: + # Multi-stage: last seq axis is the final stage + # If start token included: size == pn + 1 -> drop the first token on that axis + if hidden.shape[-2] == seq_lens[-1] + 1: + hidden = hidden[..., 1:, :] + elif hidden.shape[-2] != seq_lens[-1]: + raise AssertionError( + f"Unexpected multi-stage final seq size {hidden.shape[-2]} given seq_lens={seq_lens}" + ) + # Flatten all hierarchical sequence dims: (B, p1_prime, p2, ..., pn, D) -> (B, L_pad, D) + # We can do generic flatten: "b ... d -> b (...) d" since we only keep batch and last feature + out_hidden = rearrange(hidden, "b ... d -> b (...) d") + + # Slice to the *original* seq_len (avoid any padding mismatch) + return out_hidden[:, :seq_len, :] + + h_ids_flat = normalize_hidden(h_ids) + h_embs_flat = normalize_hidden(h_embs) + + assert ( + h_ids_flat.shape == h_embs_flat.shape + ), f"Shape mismatch after normalization: ids={h_ids_flat.shape}, embeds={h_embs_flat.shape}" + assert torch.allclose( + h_ids_flat, h_embs_flat, atol=1e-5 + ), "Hidden states differ between ids and inputs_embeds paths after normalization" diff --git a/tests/unit/model/test_multi_stage_token_embeddings.py b/tests/unit/model/test_multi_stage_token_embeddings.py new file mode 100644 index 0000000..e28f7dc --- /dev/null +++ b/tests/unit/model/test_multi_stage_token_embeddings.py @@ -0,0 +1,214 @@ +import pytest +import torch +import torch.nn as nn + +from mblm.model.multi_stage_token_embedding import MultiStageTokenEmbedding, _StageTokenEmbedding + + +@pytest.fixture(autouse=True) +def _seed(): + torch.manual_seed(1234) + + +def _build(num_tokens=257, pad_id=256, model_dims=(1024, 1024), seq_lens=(1024, 8)): + """Helper to construct the reversed module list like the legacy initializer.""" + modules = MultiStageTokenEmbedding.build( + model_dims=model_dims, + seq_lens=seq_lens, + vocab_size=num_tokens, + pad_token_id=pad_id, + ) + return modules + + +def test_returns_reversed_order_modulelist(): + model_dims = (512, 512, 512) + seq_lens = (32, 8, 4) + modules = _build(model_dims=model_dims, seq_lens=seq_lens) + # Expect one local stage + len(model_dims)-1 global stages + assert isinstance(modules, nn.ModuleList) + assert len(modules) == len(model_dims) # local + all globals + # First element must be local stage module (last in the hierarchy) + assert isinstance(modules[0], _StageTokenEmbedding) + assert modules[0].is_local is True + # The rest are global + for module in modules[1:]: + assert isinstance(module, _StageTokenEmbedding) + assert module.is_local is False + + +def test_local_stage_ids_vs_embeds_identical(): + num_tokens, pad_id = 260, 0 + model_dims = (128, 128) + seq_lens = (16, 4) + modules = _build(num_tokens=num_tokens, pad_id=pad_id, model_dims=model_dims, seq_lens=seq_lens) + + # local stage is index 0 in the reversed list + local_stage = modules[0] + batch_size, seq_len = 2, 11 + input_ids = torch.randint(0, num_tokens, (batch_size, seq_len)) + + # Path 1: via input_ids + out_ids = local_stage(input_ids=input_ids) + + # Path 2: via inputs_embeds (same embedding lookup externally) + embeds = local_stage.embedding(input_ids) + out_embeds = local_stage(inputs_embeds=embeds) + + assert out_ids.shape == (batch_size, seq_len, model_dims[-1]) + assert torch.allclose(out_ids, out_embeds, atol=0, rtol=0) + + +def test_padding_row_is_zero_for_local_stage(): + num_tokens, pad_id = 260, 5 + model_dims = (64, 64) + seq_lens = (8, 2) + modules = _build(num_tokens=num_tokens, pad_id=pad_id, model_dims=model_dims, seq_lens=seq_lens) + + local_stage = modules[0] + with torch.no_grad(): + # padding row must be zeros + row = local_stage.embedding.weight[pad_id] + assert torch.all(row == 0) + + +def test_error_if_both_or_none_inputs(): + modules = _build() + local_stage = modules[0] + + batch_size, seq_len = 2, 7 + ids = torch.randint(0, 257, (batch_size, seq_len)) + embeds = local_stage.embedding(ids) + + with pytest.raises(ValueError): + _ = local_stage() # neither + + with pytest.raises(ValueError): + _ = local_stage(input_ids=ids, inputs_embeds=embeds) # both + + +def test_global_stage_ids_vs_embeds_identical_projection(): + """ + For a global stage: + - ids path: ids -> embedding -> [B, R, d] -> flatten/proj -> [B, D_model] + - embeds path: [B, R, d] -> flatten/proj -> [B, D_model] + If embeds come from the SAME embedding lookup, outputs must match exactly. + """ + num_tokens, pad_id = 257, 256 + # 3 stages: D1(global), D2(global), D3(local=last) + model_dims = (384, 256, 128) + seq_lens = (16, 4, 2) # P2, P3 used in global stages + modules = _build(num_tokens=num_tokens, pad_id=pad_id, model_dims=model_dims, seq_lens=seq_lens) + + # modules[0] -> local (D3=128) + # modules[1] -> global for D2=256 with patch_size = P3=2 + # modules[2] -> global for D1=384 with patch_size = P2*P3 + global_stage = modules[1] # the one right after local + + # R is patch_size for this global stage: + patch_size = global_stage.patch_size + batch_size = 3 + + # ids path + ids = torch.randint(0, num_tokens, (batch_size, patch_size)) + out_ids = global_stage(input_ids=ids) + + # embeds path, but ensure they come from the SAME embedding weights + embeds = global_stage.embedding(ids) # [B, R, d] + out_embeds = global_stage(inputs_embeds=embeds) + + assert out_ids.shape == (batch_size, model_dims[-2]) # D2 + assert torch.allclose(out_ids, out_embeds, atol=0, rtol=0) + + +def test_global_patch_size_accumulation_matches_seq_lens(): + """ + For n stages, reversed globals should have patch sizes: + stage n-1: prod(seq_lens[n]) (= P_n) + stage n-2: prod(seq_lens[n-1:n]) (= P_{n-1} * P_n) + ... + """ + model_dims = (256, 192, 128, 96) # 4 stages, last=local + seq_lens = (8, 4, 2, 1) + modules = _build(model_dims=model_dims, seq_lens=seq_lens) + + # modules: [local, global_3, global_2, global_1] + local_stage, global_stage_3, global_stage_2, global_stage_1 = modules + + assert local_stage.is_local + # Expected patch sizes: + p2, p3, p4 = seq_lens[1], seq_lens[2], seq_lens[3] + assert global_stage_3.patch_size == p4 + assert global_stage_2.patch_size == p3 * p4 + assert global_stage_1.patch_size == p2 * p3 * p4 + + +def test_local_shapes_and_types(): + modules = _build() + local_stage = modules[0] + + batch_size, seq_len = 2, 9 + ids = torch.randint(0, 257, (batch_size, seq_len)) + out = local_stage(input_ids=ids) + + assert out.dtype == local_stage.embedding.weight.dtype + assert out.shape[-1] == local_stage.local_dim + + +def test_global_shapes_and_types(): + modules = _build(model_dims=(256, 256, 256), seq_lens=(12, 3, 2)) + # global stage right after local + global_stage = modules[1] + + batch_size = 4 + patch_size = global_stage.patch_size # 2 + ids = torch.randint(0, 257, (batch_size, patch_size)) + out = global_stage(input_ids=ids) + + assert out.dtype == global_stage.embedding.weight.dtype + assert out.shape == (batch_size, 256) + + +def test_gradient_flows_through_both_paths(): + """ + Ensure we can backprop through ids->embed->proj and embeds->proj paths. + """ + num_tokens = 300 + modules = _build(num_tokens=num_tokens, model_dims=(128, 128, 128), seq_lens=(10, 5, 2)) + local_stage = modules[0] + global_stage = modules[1] + + # Local: ids path + ids_local = torch.randint(0, num_tokens, (2, 7)) + out_local = local_stage(input_ids=ids_local) + loss_local = out_local.pow(2).mean() + loss_local.backward(retain_graph=True) + # At least embedding grads for some rows should exist (non-padding) + assert local_stage.embedding.weight.grad is not None + + # Global: embeds path + patch_size = global_stage.patch_size + ids_global = torch.randint(0, num_tokens, (2, patch_size)) + embeds = global_stage.embedding(ids_global).detach().requires_grad_(True) + out_global = global_stage(inputs_embeds=embeds) + loss_global = out_global.abs().mean() + loss_global.backward() + # Check that projection (Linear) received gradients + linear_layer = next(m for m in global_stage._post if isinstance(m, nn.Linear)) + assert linear_layer.weight.grad is not None + + +def test_inputs_embeds_must_have_correct_shape_for_local_and_global(): + modules = _build(model_dims=(128, 96, 64), seq_lens=(8, 4, 2)) + local_stage = modules[0] + global_stage = modules[1] # has patch_size = 2 + + # local expects [..., L, d] + wrong_local_embeds = torch.randn(2, 5, local_stage.local_dim + 1) + with pytest.raises(RuntimeError): + _ = local_stage(inputs_embeds=wrong_local_embeds) + + # global expects [..., R, d] + wrong_global_embeds = torch.randn(2, global_stage.patch_size + 1, global_stage.local_dim) + with pytest.raises(RuntimeError): + _ = global_stage(inputs_embeds=wrong_global_embeds) diff --git a/tests/unit/utils/test_io.py b/tests/unit/utils/test_io.py index c3c660a..57d6c98 100644 --- a/tests/unit/utils/test_io.py +++ b/tests/unit/utils/test_io.py @@ -5,14 +5,16 @@ from concurrent.futures import ThreadPoolExecutor from datetime import datetime from pathlib import Path -from typing import NamedTuple, cast +from typing import NamedTuple import pytest import torch from pydantic import BaseModel +from torch import nn from mblm import MBLM, MBLMModelConfig, TransformerBlock from mblm.model.embeddings import MBLM_TOKEN_EMB_MIGRATION +from mblm.model.multi_stage_token_embedding import _StageTokenEmbedding from mblm.utils.io import ( CSVWriter, NDJSONWriter, @@ -169,8 +171,17 @@ def create_model(num_tokens: int): ) num_src_emb, num_tgt_emb = 5, 6 + pad_id = 0 + model_src = create_model(num_src_emb) model_tgt = create_model(num_tgt_emb) + + # Assert updated structure + assert isinstance(model_src.token_embs_rev[0], _StageTokenEmbedding) + assert isinstance(model_src.token_embs_rev[1], _StageTokenEmbedding) + assert isinstance(model_tgt.token_embs_rev[0], _StageTokenEmbedding) + assert isinstance(model_tgt.token_embs_rev[1], _StageTokenEmbedding) + with tempfile.TemporaryDirectory() as tmpdir: _, chkpoint = save_model_state(tmpdir, "checkpoint", model_src, 0) model_tgt, _ = load_model_state( @@ -178,48 +189,55 @@ def create_model(num_tokens: int): model_tgt, map_extend_embeddings=MBLM_TOKEN_EMB_MIGRATION, ) - """ - This is the structure of the embeddings we're migrating: - - (token_embs_rev): ModuleList( - case 1: (0): Embedding(255, 512, padding_idx=0) - (1): Sequential( - case 2: (0): Embedding(255, 512, padding_idx=0) - (1): Rearrange('... r d -> ... (r d)') - (2): LayerNorm((4096,), eps=1e-05, elementwise_affine=True) - (3): Linear(in_features=4096, out_features=1024, bias=True) - (4): LayerNorm((1024,), eps=1e-05, elementwise_affine=True) - ) - ... - cases 3/4: (to_logits): Linear(in_features=512, out_features=255, bias=True) - ) - """ - src_emb = cast(torch.nn.Embedding, model_src.token_embs_rev[0]) - src_emb_seq = cast(torch.nn.Sequential, model_src.token_embs_rev[1]) - tgt_emb = cast(torch.nn.Embedding, model_tgt.token_embs_rev[0]) - tgt_emb_seq = cast(torch.nn.Sequential, model_tgt.token_embs_rev[1]) - - # case 1, base embedding - assert tgt_emb.num_embeddings == num_tgt_emb - assert tgt_emb.weight[:num_src_emb].equal(src_emb.weight) - # case 2, embedding in sequential - assert tgt_emb_seq[0].num_embeddings == num_tgt_emb - assert tgt_emb_seq[0].weight[:num_src_emb].equal(src_emb_seq[0].weight) - # case 3/4, logits + # Extract embeddings (new layout: embedding lives at .embedding) + src_stage0_emb: nn.Embedding = model_src.token_embs_rev[0].embedding + src_stage1_emb: nn.Embedding = model_src.token_embs_rev[1].embedding + tgt_stage0_emb: nn.Embedding = model_tgt.token_embs_rev[0].embedding + tgt_stage1_emb: nn.Embedding = model_tgt.token_embs_rev[1].embedding + + # Sizes grew + assert tgt_stage0_emb.num_embeddings == num_tgt_emb + assert tgt_stage1_emb.num_embeddings == num_tgt_emb assert model_tgt.to_logits.weight.size(0) == num_tgt_emb assert model_tgt.to_logits.bias.size(0) == num_tgt_emb - assert model_tgt.to_logits.weight[:num_src_emb].equal(model_src.to_logits.weight) - assert model_tgt.to_logits.bias[:num_src_emb].equal(model_src.to_logits.bias) - # check if new token id works + # === Functional equality for old token ids (skip pad if migration handles it specially) === + # Compare the result of embedding lookups rather than raw weight slicing. + old_token_ids = torch.arange(num_src_emb, dtype=torch.long) + non_pad_ids = old_token_ids[old_token_ids != pad_id] + assert non_pad_ids.numel() > 0, "Expected at least one non-pad id in source vocab" + + # Stage 0 (local) functional check + src_s0_vecs = src_stage0_emb(non_pad_ids) # [K, D0] + tgt_s0_vecs = tgt_stage0_emb(non_pad_ids) # [K, D0] + assert torch.allclose( + tgt_s0_vecs, src_s0_vecs, atol=0, rtol=0 + ), "Stage-0 embeddings differ for existing token ids" + + # Stage 1 (global) functional check (embedding part only) + src_s1_vecs = src_stage1_emb(non_pad_ids) # [K, D1] + tgt_s1_vecs = tgt_stage1_emb(non_pad_ids) # [K, D1] + assert torch.allclose( + tgt_s1_vecs, src_s1_vecs, atol=0, rtol=0 + ), "Stage-1 embeddings differ for existing token ids" + + # === Logits preservation for old ids === + assert torch.allclose( + model_tgt.to_logits.weight[:num_src_emb], model_src.to_logits.weight, atol=0, rtol=0 + ), "to_logits.weight rows for old ids not preserved" + assert torch.allclose( + model_tgt.to_logits.bias[:num_src_emb], model_src.to_logits.bias, atol=0, rtol=0 + ), "to_logits.bias rows for old ids not preserved" + + # === New token id should be accepted only by the migrated model === max_new_token_id = num_tgt_emb - 1 - input_for_tgt_model_only = torch.tensor([[max_new_token_id]]).long() + input_for_tgt_model_only = torch.tensor([[max_new_token_id]], dtype=torch.long) + with pytest.raises(Exception): - # should fail for old model model_src.forward(input_for_tgt_model_only) + try: - # should work for migrated model model_tgt.forward(input_for_tgt_model_only) except Exception as error: pytest.fail(f"Forward pass should work: {error}")