Add files using upload-large-folder tool

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/markdown_it/_punycode.py

@@ -0,0 +1,67 @@
+# Copyright 2014 Mathias Bynens <https://mathiasbynens.be/>
+# Copyright 2021 Taneli Hukkinen
+#
+# 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.
+
+import codecs
+from collections.abc import Callable
+import re
+
+REGEX_SEPARATORS = re.compile(r"[\x2E\u3002\uFF0E\uFF61]")
+REGEX_NON_ASCII = re.compile(r"[^\0-\x7E]")
+
+
+def encode(uni: str) -> str:
+    return codecs.encode(uni, encoding="punycode").decode()
+
+
+def decode(ascii: str) -> str:
+    return codecs.decode(ascii, encoding="punycode")  # type: ignore
+
+
+def map_domain(string: str, fn: Callable[[str], str]) -> str:
+    parts = string.split("@")
+    result = ""
+    if len(parts) > 1:
+        # In email addresses, only the domain name should be punycoded. Leave
+        # the local part (i.e. everything up to `@`) intact.
+        result = parts[0] + "@"
+        string = parts[1]
+    labels = REGEX_SEPARATORS.split(string)
+    encoded = ".".join(fn(label) for label in labels)
+    return result + encoded
+
+
+def to_unicode(obj: str) -> str:
+    def mapping(obj: str) -> str:
+        if obj.startswith("xn--"):
+            return decode(obj[4:].lower())
+        return obj
+
+    return map_domain(obj, mapping)
+
+
+def to_ascii(obj: str) -> str:
+    def mapping(obj: str) -> str:
+        if REGEX_NON_ASCII.search(obj):
+            return "xn--" + encode(obj)
+        return obj
+
+    return map_domain(obj, mapping)

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/markdown_it/port.yaml

@@ -0,0 +1,48 @@
+- package: markdown-it/markdown-it
+  version: 14.1.0
+  commit: 0fe7ccb4b7f30236fb05f623be6924961d296d3d
+  date: Mar 19, 2024
+  notes:
+    - Rename variables that use python built-in names, e.g.
+      - `max` -> `maximum`
+      - `len` -> `length`
+      - `str` -> `string`
+    - |
+      Convert JS `for` loops to `while` loops
+      this is generally the main difference between the codes,
+      because in python you can't do e.g. `for {i=1;i<x;i++} {}`
+    - |
+      `env` is a common Python dictionary, and so does not have attribute access to keys,
+      as with JavaScript dictionaries.
+      `options` have attribute access only to core markdownit configuration options
+    - |
+      `Token.attrs` is a dictionary, instead of a list of lists.
+      Upstream the list format is only used to guarantee order: https://github.com/markdown-it/markdown-it/issues/142,
+      but in Python 3.7+ order of dictionaries is guaranteed.
+      One should anyhow use the `attrGet`, `attrSet`, `attrPush` and `attrJoin` methods
+      to manipulate `Token.attrs`, which have an identical signature to those upstream.
+    - Use python version of `charCodeAt`
+    - |
+      Use `str` units instead of `int`s to represent Unicode codepoints.
+      This provides a significant performance boost
+    - |
+      In markdown_it/rules_block/reference.py,
+      record line range in state.env["references"] and add state.env["duplicate_refs"]
+      This is to allow renderers to report on issues regarding references
+    - |
+      The `MarkdownIt.__init__` signature is slightly different for updating options,
+      since you must always specify the config first, e.g.
+      use `MarkdownIt("commonmark", {"html": False})` instead of `MarkdownIt({"html": False})`
+    - The default configuration preset for `MarkdownIt` is "commonmark" not "default"
+    - Allow custom renderer to be passed to `MarkdownIt`
+    - |
+      change render method signatures
+      `func(tokens, idx, options, env, slf)` to
+      `func(self, tokens, idx, options, env)`
+    - |
+      Extensions add render methods by format
+      `MarkdownIt.add_render_rule(name, function, fmt="html")`,
+      rather than `MarkdownIt.renderer.rules[name] = function`
+      and renderers should declare a class property `__output__ = "html"`.
+      This allows for extensibility to more than just HTML renderers
+    - inline tokens in tables are assigned a map (this is helpful for propagation to children)

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/byt5/tokenization_byt5.py

@@ -0,0 +1,234 @@
+# Copyright 2021 T5 Authors and HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tokenization class for model ByT5."""
+
+import warnings
+
+from ...tokenization_python import AddedToken, PreTrainedTokenizer
+from ...utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class ByT5Tokenizer(PreTrainedTokenizer):
+    """
+    Construct a ByT5 tokenizer. ByT5 simply uses raw bytes utf-8 encoding.
+
+    This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to
+    this superclass for more information regarding those methods.
+
+    Args:
+        eos_token (`str`, *optional*, defaults to `"</s>"`):
+            The end of sequence token.
+
+            <Tip>
+
+            When building a sequence using special tokens, this is not the token that is used for the end of sequence.
+            The token used is the `sep_token`.
+
+            </Tip>
+
+        unk_token (`str`, *optional*, defaults to `"<unk>"`):
+            The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
+            token instead.
+        pad_token (`str`, *optional*, defaults to `"<pad>"`):
+            The token used for padding, for example when batching sequences of different lengths.
+        extra_ids (`int`, *optional*, defaults to 125):
+            Add a number of extra ids added to the end of the vocabulary for use as sentinels. These tokens are
+            accessible as "<extra_id_{%d}>" where "{%d}" is a number between 0 and extra_ids-1. Extra tokens are
+            indexed from the end of the vocabulary up to beginning ("<extra_id_0>" is the last token in the vocabulary
+            like in ByT5 preprocessing see
+            [here](https://github.com/google-research/text-to-text-transfer-transformer/blob/9fd7b14a769417be33bc6c850f9598764913c833/t5/data/preprocessors.py#L2117)).
+        additional_special_tokens (`list[str]`, *optional*):
+            Additional special tokens used by the tokenizer.
+    """
+
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(
+        self,
+        eos_token="</s>",
+        unk_token="<unk>",
+        pad_token="<pad>",
+        extra_ids=125,
+        additional_special_tokens=None,
+        **kwargs,
+    ) -> None:
+        # Add extra_ids to the special token list
+        if extra_ids > 0 and additional_special_tokens is None:
+            additional_special_tokens = [f"<extra_id_{i}>" for i in range(extra_ids)]
+        elif extra_ids > 0 and additional_special_tokens is not None and len(additional_special_tokens) > 0:
+            # Check that we have the right number of extra_id special tokens
+            extra_tokens = len(set(filter(lambda x: bool("extra_id" in str(x)), additional_special_tokens)))
+            if extra_tokens != extra_ids:
+                raise ValueError(
+                    f"Both extra_ids ({extra_ids}) and additional_special_tokens ({additional_special_tokens}) are"
+                    " provided to ByT5Tokenizer. In this case the additional_special_tokens must include the"
+                    " extra_ids tokens"
+                )
+
+        pad_token = AddedToken(pad_token, lstrip=True, rstrip=True) if isinstance(pad_token, str) else pad_token
+        # we force left and right stripping for backward compatibility. The byt5tests depend on this.
+        eos_token = AddedToken(eos_token, lstrip=True, rstrip=True) if isinstance(eos_token, str) else eos_token
+        unk_token = AddedToken(unk_token, lstrip=True, rstrip=True) if isinstance(unk_token, str) else unk_token
+        # unk token needs to be in the vocab with correct index
+        self._added_tokens_decoder = {0: pad_token, 1: eos_token, 2: unk_token}
+        self.offset = len(self._added_tokens_decoder)
+        self._utf_vocab_size = 2**8  # utf is 8 bits
+        super().__init__(
+            eos_token=eos_token,
+            unk_token=unk_token,
+            pad_token=pad_token,
+            extra_ids=0,
+            additional_special_tokens=additional_special_tokens,  # TODO extra ids are not used :sweatywmile:
+            **kwargs,
+        )
+
+    @property
+    def vocab_size(self):
+        return self._utf_vocab_size
+
+    def get_vocab(self):
+        vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size + self.offset)}
+        vocab.update(self.added_tokens_encoder)
+        return vocab
+
+    def get_special_tokens_mask(
+        self, token_ids_0: list[int], token_ids_1: list[int] | None = None, already_has_special_tokens: bool = False
+    ) -> list[int]:
+        """
+        Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
+        special tokens using the tokenizer `prepare_for_model` method.
+
+        Args:
+            token_ids_0 (`list[int]`):
+                List of IDs.
+            token_ids_1 (`list[int]`, *optional*):
+                Optional second list of IDs for sequence pairs.
+            already_has_special_tokens (`bool`, *optional*, defaults to `False`):
+                Whether or not the token list is already formatted with special tokens for the model.
+
+        Returns:
+            `list[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
+        """
+        if already_has_special_tokens:
+            return super().get_special_tokens_mask(
+                token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
+            )
+
+        # normal case: some special tokens
+        if token_ids_1 is None:
+            return ([0] * len(token_ids_0)) + [1]
+        return ([0] * len(token_ids_0)) + [1] + ([0] * len(token_ids_1)) + [1]
+
+    def _add_eos_if_not_present(self, token_ids: list[int]) -> list[int]:
+        """Do not add eos again if user already added it."""
+        if len(token_ids) > 0 and token_ids[-1] == self.eos_token_id:
+            warnings.warn(
+                f"This sequence already has {self.eos_token}. In future versions this behavior may lead to duplicated"
+                " eos tokens being added."
+            )
+            return token_ids
+        else:
+            return token_ids + [self.eos_token_id]
+
+    def create_token_type_ids_from_sequences(
+        self, token_ids_0: list[int], token_ids_1: list[int] | None = None
+    ) -> list[int]:
+        """
+        Create a mask from the two sequences passed to be used in a sequence-pair classification task. ByT5 does not
+        make use of token type ids, therefore a list of zeros is returned.
+
+        Args:
+            token_ids_0 (`list[int]`):
+                List of IDs.
+            token_ids_1 (`list[int]`, *optional*):
+                Optional second list of IDs for sequence pairs.
+
+        Returns:
+            `list[int]`: List of zeros.
+        """
+        eos = [self.eos_token_id]
+
+        if token_ids_1 is None:
+            return len(token_ids_0 + eos) * [0]
+        return len(token_ids_0 + eos + token_ids_1 + eos) * [0]
+
+    def build_inputs_with_special_tokens(
+        self, token_ids_0: list[int], token_ids_1: list[int] | None = None
+    ) -> list[int]:
+        """
+        Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
+        adding special tokens. A sequence has the following format:
+
+        - single sequence: `X </s>`
+        - pair of sequences: `A </s> B </s>`
+
+        Args:
+            token_ids_0 (`list[int]`):
+                List of IDs to which the special tokens will be added.
+            token_ids_1 (`list[int]`, *optional*):
+                Optional second list of IDs for sequence pairs.
+
+        Returns:
+            `list[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
+        """
+        token_ids_0 = self._add_eos_if_not_present(token_ids_0)
+        if token_ids_1 is None:
+            return token_ids_0
+        else:
+            token_ids_1 = self._add_eos_if_not_present(token_ids_1)
+            return token_ids_0 + token_ids_1
+
+    def _tokenize(self, text: str) -> list[str]:
+        """Take as input a string and return a list of strings (tokens) for words/sub-words"""
+        tokens = [chr(i) for i in text.encode("utf-8")]
+        return tokens
+
+    def _convert_token_to_id(self, token):
+        """Converts a token (str) in an id using the vocab."""
+
+        if len(token) != 1:
+            token_id = None
+        else:
+            token_id = ord(token) + self.offset
+
+        return token_id
+
+    def _convert_id_to_token(self, index):
+        """Converts an index (integer) in a token (str) using the vocab."""
+        token = chr(index - self.offset)
+        return token
+
+    def convert_tokens_to_string(self, tokens):
+        """Converts a sequence of tokens (string) in a single string."""
+        bstring = b""
+        for token in tokens:
+            if token in self.added_tokens_decoder:
+                tok_string = self.added_tokens_decoder[token].encode("utf-8")
+            elif token in self.added_tokens_encoder:
+                tok_string = token.encode("utf-8")
+            else:
+                tok_string = bytes([ord(token)])
+            bstring += tok_string
+        string = bstring.decode("utf-8", errors="ignore")
+        return string
+
+    # ByT5Tokenizer has no vocab file
+    def save_vocabulary(self, save_directory: str, filename_prefix: str | None = None) -> tuple[str]:
+        return ()
+
+
+__all__ = ["ByT5Tokenizer"]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/convnextv2/__init__.py

@@ -0,0 +1,27 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+    from .configuration_convnextv2 import *
+    from .modeling_convnextv2 import *
+else:
+    import sys
+
+    _file = globals()["__file__"]
+    sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/convnextv2/modeling_convnextv2.py

@@ -0,0 +1,428 @@
+# Copyright 2023 Meta Platforms, Inc. and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PyTorch ConvNextV2 model."""
+
+import torch
+from torch import nn
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...backbone_utils import BackboneMixin, filter_output_hidden_states
+from ...modeling_outputs import (
+    BackboneOutput,
+    BaseModelOutputWithNoAttention,
+    BaseModelOutputWithPoolingAndNoAttention,
+    ImageClassifierOutputWithNoAttention,
+)
+from ...modeling_utils import PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring, logging
+from ...utils.generic import can_return_tuple, merge_with_config_defaults
+from ...utils.output_capturing import capture_outputs
+from .configuration_convnextv2 import ConvNextV2Config
+
+
+logger = logging.get_logger(__name__)
+
+
+class ConvNextV2GRN(nn.Module):
+    """GRN (Global Response Normalization) layer"""
+
+    def __init__(self, dim: int):
+        super().__init__()
+        self.weight = nn.Parameter(torch.zeros(1, 1, 1, dim))
+        self.bias = nn.Parameter(torch.zeros(1, 1, 1, dim))
+
+    def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor:
+        # Compute and normalize global spatial feature maps
+        global_features = torch.linalg.vector_norm(hidden_states, ord=2, dim=(1, 2), keepdim=True)
+        norm_features = global_features / (global_features.mean(dim=-1, keepdim=True) + 1e-6)
+        hidden_states = self.weight * (hidden_states * norm_features) + self.bias + hidden_states
+
+        return hidden_states
+
+
+# Copied from transformers.models.convnext.modeling_convnext.ConvNextLayerNorm with ConvNext->ConvNextV2
+class ConvNextV2LayerNorm(nn.LayerNorm):
+    r"""LayerNorm that supports two data formats: channels_last (default) or channels_first.
+    The ordering of the dimensions in the inputs. channels_last corresponds to inputs with shape (batch_size, height,
+    width, channels) while channels_first corresponds to inputs with shape (batch_size, channels, height, width).
+    """
+
+    def __init__(self, normalized_shape, *, eps=1e-6, data_format="channels_last", **kwargs):
+        super().__init__(normalized_shape, eps=eps, **kwargs)
+        if data_format not in ["channels_last", "channels_first"]:
+            raise NotImplementedError(f"Unsupported data format: {data_format}")
+        self.data_format = data_format
+
+    def forward(self, features: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            features: Tensor of shape (batch_size, channels, height, width) OR (batch_size, height, width, channels)
+        """
+        if self.data_format == "channels_first":
+            features = features.permute(0, 2, 3, 1)
+            features = super().forward(features)
+            features = features.permute(0, 3, 1, 2)
+        else:
+            features = super().forward(features)
+        return features
+
+
+# Copied from transformers.models.convnext.modeling_convnext.ConvNextEmbeddings with ConvNext->ConvNextV2
+class ConvNextV2Embeddings(nn.Module):
+    """This class is comparable to (and inspired by) the SwinEmbeddings class
+    found in src/transformers/models/swin/modeling_swin.py.
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        self.patch_embeddings = nn.Conv2d(
+            config.num_channels, config.hidden_sizes[0], kernel_size=config.patch_size, stride=config.patch_size
+        )
+        self.layernorm = ConvNextV2LayerNorm(config.hidden_sizes[0], eps=1e-6, data_format="channels_first")
+        self.num_channels = config.num_channels
+
+    def forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor:
+        num_channels = pixel_values.shape[1]
+        if num_channels != self.num_channels:
+            raise ValueError(
+                "Make sure that the channel dimension of the pixel values match with the one set in the configuration."
+            )
+        embeddings = self.patch_embeddings(pixel_values)
+        embeddings = self.layernorm(embeddings)
+        return embeddings
+
+
+# Copied from transformers.models.swin.modular_swin.SwinDropPath with SwinDropPath->ConvNextV2DropPath
+class ConvNextV2DropPath(nn.Module):
+    """Stochastic depth (DropPath) per sample, for residual blocks.
+
+    Identity when ``drop_prob`` is 0 or outside training. See `Deep Networks with Stochastic Depth
+    <https://arxiv.org/abs/1603.09382>`_.
+    """
+
+    def __init__(self, drop_prob: float = 0.0) -> None:
+        super().__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        if self.drop_prob == 0.0 or not self.training:
+            return hidden_states
+        keep_prob = 1 - self.drop_prob
+        shape = (hidden_states.shape[0],) + (1,) * (hidden_states.ndim - 1)
+        random_tensor = torch.rand(shape, dtype=hidden_states.dtype, device=hidden_states.device)
+        random_tensor = torch.floor(random_tensor + keep_prob)
+        return hidden_states.div(keep_prob) * random_tensor
+
+    def extra_repr(self) -> str:
+        return f"p={self.drop_prob}"
+
+
+class ConvNextV2Layer(nn.Module):
+    """This corresponds to the `Block` class in the original implementation.
+
+    There are two equivalent implementations: [DwConv, LayerNorm (channels_first), Conv, GELU,1x1 Conv]; all in (N, C,
+    H, W) (2) [DwConv, Permute to (N, H, W, C), LayerNorm (channels_last), Linear, GELU, Linear]; Permute back
+
+    The authors used (2) as they find it slightly faster in PyTorch.
+
+    Args:
+        config ([`ConvNextV2Config`]): Model configuration class.
+        dim (`int`): Number of input channels.
+        drop_path (`float`): Stochastic depth rate. Default: 0.0.
+    """
+
+    def __init__(self, config, dim, drop_path=0):
+        super().__init__()
+        # depthwise conv
+        self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim)
+        self.layernorm = ConvNextV2LayerNorm(dim, eps=1e-6)
+        # pointwise/1x1 convs, implemented with linear layers
+        self.pwconv1 = nn.Linear(dim, 4 * dim)
+        self.act = ACT2FN[config.hidden_act]
+        self.grn = ConvNextV2GRN(4 * dim)
+        self.pwconv2 = nn.Linear(4 * dim, dim)
+        self.drop_path = ConvNextV2DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+    def forward(self, features: torch.Tensor) -> torch.Tensor:
+        residual = features
+        features = self.dwconv(features)
+        # (batch_size, num_channels, height, width) -> (batch_size, height, width, num_channels)
+        features = features.permute(0, 2, 3, 1)
+        features = self.layernorm(features)
+        features = self.pwconv1(features)
+        features = self.act(features)
+        features = self.grn(features)
+        features = self.pwconv2(features)
+        # (batch_size, height, width, num_channels) -> (batch_size, num_channels, height, width)
+        features = features.permute(0, 3, 1, 2)
+
+        features = residual + self.drop_path(features)
+        return features
+
+
+# Copied from transformers.models.convnext.modeling_convnext.ConvNextStage with ConvNeXT->ConvNeXTV2, ConvNext->ConvNextV2
+class ConvNextV2Stage(nn.Module):
+    """ConvNeXTV2 stage, consisting of an optional downsampling layer + multiple residual blocks.
+
+    Args:
+        config ([`ConvNextV2Config`]): Model configuration class.
+        in_channels (`int`): Number of input channels.
+        out_channels (`int`): Number of output channels.
+        depth (`int`): Number of residual blocks.
+        drop_path_rates(`list[float]`): Stochastic depth rates for each layer.
+    """
+
+    def __init__(self, config, in_channels, out_channels, kernel_size=2, stride=2, depth=2, drop_path_rates=None):
+        super().__init__()
+
+        if in_channels != out_channels or stride > 1:
+            self.downsampling_layer = nn.ModuleList(
+                [
+                    ConvNextV2LayerNorm(in_channels, eps=1e-6, data_format="channels_first"),
+                    nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride),
+                ]
+            )
+        else:
+            self.downsampling_layer = nn.ModuleList()
+        drop_path_rates = drop_path_rates or [0.0] * depth
+        self.layers = nn.ModuleList(
+            [ConvNextV2Layer(config, dim=out_channels, drop_path=drop_path_rates[j]) for j in range(depth)]
+        )
+
+    def forward(self, features: torch.Tensor) -> torch.Tensor:
+        for layer in self.downsampling_layer:
+            features = layer(features)
+        for layer in self.layers:
+            features = layer(features)
+        return features
+
+
+@auto_docstring
+class ConvNextV2PreTrainedModel(PreTrainedModel):
+    config: ConvNextV2Config
+    base_model_prefix = "convnextv2"
+    main_input_name = "pixel_values"
+    input_modalities = ("image",)
+    _no_split_modules = ["ConvNextV2Layer"]
+
+    @torch.no_grad()
+    def _init_weights(self, module):
+        """Initialize the weights"""
+        super()._init_weights(module)
+        if isinstance(module, ConvNextV2GRN):
+            init.zeros_(module.weight)
+            init.zeros_(module.bias)
+
+
+# Copied from transformers.models.convnext.modeling_convnext.ConvNextEncoder with CONVNEXT->CONVNEXTV2, ConvNext->ConvNextV2
+class ConvNextV2Encoder(ConvNextV2PreTrainedModel):
+    main_input_name = "hidden_states"
+    _can_record_outputs = {"hidden_states": ConvNextV2Stage}
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.stages = nn.ModuleList()
+        drop_path_rates = [
+            x.tolist()
+            for x in torch.linspace(0, config.drop_path_rate, sum(config.depths), device="cpu").split(config.depths)
+        ]
+        prev_chs = config.hidden_sizes[0]
+        for i in range(config.num_stages):
+            out_chs = config.hidden_sizes[i]
+            stage = ConvNextV2Stage(
+                config,
+                in_channels=prev_chs,
+                out_channels=out_chs,
+                stride=2 if i > 0 else 1,
+                depth=config.depths[i],
+                drop_path_rates=drop_path_rates[i],
+            )
+            self.stages.append(stage)
+            prev_chs = out_chs
+
+        self.post_init()
+
+    @merge_with_config_defaults
+    @capture_outputs(tie_last_hidden_states=False)
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BaseModelOutputWithNoAttention:
+        for layer_module in self.stages:
+            hidden_states = layer_module(hidden_states)
+
+        return BaseModelOutputWithNoAttention(last_hidden_state=hidden_states)
+
+
+@auto_docstring
+# Copied from transformers.models.convnext.modeling_convnext.ConvNextModel with CONVNEXT->CONVNEXTV2, ConvNext->ConvNextV2
+class ConvNextV2Model(ConvNextV2PreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+
+        self.embeddings = ConvNextV2Embeddings(config)
+        self.encoder = ConvNextV2Encoder(config)
+
+        # final layernorm layer
+        self.layernorm = nn.LayerNorm(config.hidden_sizes[-1], eps=config.layer_norm_eps)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @can_return_tuple
+    @auto_docstring
+    def forward(
+        self, pixel_values: torch.FloatTensor | None = None, **kwargs: Unpack[TransformersKwargs]
+    ) -> BaseModelOutputWithPoolingAndNoAttention:
+        if pixel_values is None:
+            raise ValueError("You have to specify pixel_values")
+
+        embedding_output = self.embeddings(pixel_values)
+        encoder_outputs: BaseModelOutputWithNoAttention = self.encoder(embedding_output, **kwargs)
+        last_hidden_state = encoder_outputs.last_hidden_state
+
+        # global average pooling, (N, C, H, W) -> (N, C)
+        pooled_output = self.layernorm(last_hidden_state.mean([-2, -1]))
+
+        return BaseModelOutputWithPoolingAndNoAttention(
+            last_hidden_state=last_hidden_state,
+            pooler_output=pooled_output,
+            hidden_states=encoder_outputs.hidden_states,
+        )
+
+
+@auto_docstring(
+    custom_intro="""
+    ConvNextV2 Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for
+    ImageNet.
+    """
+)
+# Copied from transformers.models.convnext.modeling_convnext.ConvNextForImageClassification with CONVNEXT->CONVNEXTV2,ConvNext->ConvNextV2,convnext->convnextv2
+class ConvNextV2ForImageClassification(ConvNextV2PreTrainedModel):
+    accepts_loss_kwargs = False
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.num_labels = config.num_labels
+        self.convnextv2 = ConvNextV2Model(config)
+
+        # Classifier head
+        if config.num_labels > 0:
+            self.classifier = nn.Linear(config.hidden_sizes[-1], config.num_labels)
+        else:
+            self.classifier = nn.Identity()
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @can_return_tuple
+    @auto_docstring
+    def forward(
+        self, pixel_values: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, **kwargs
+    ) -> ImageClassifierOutputWithNoAttention:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Labels for computing the image classification/regression loss. Indices should be in `[0, ...,
+            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+        """
+        outputs: BaseModelOutputWithPoolingAndNoAttention = self.convnextv2(pixel_values, **kwargs)
+        pooled_output = outputs.pooler_output
+        logits = self.classifier(pooled_output)
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(labels=labels, pooled_logits=logits, config=self.config)
+
+        return ImageClassifierOutputWithNoAttention(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+        )
+
+
+@auto_docstring(
+    custom_intro="""
+    ConvNeXT V2 backbone, to be used with frameworks like DETR and MaskFormer.
+    """
+)
+# Copied from transformers.models.convnext.modeling_convnext.ConvNextBackbone with CONVNEXT->CONVNEXTV2,ConvNext->ConvNextV2,facebook/convnext-tiny-224->facebook/convnextv2-tiny-1k-224
+class ConvNextV2Backbone(BackboneMixin, ConvNextV2PreTrainedModel):
+    has_attentions = False
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.embeddings = ConvNextV2Embeddings(config)
+        self.encoder = ConvNextV2Encoder(config)
+        self.num_features = [config.hidden_sizes[0]] + config.hidden_sizes
+
+        # Add layer norms to hidden states of out_features
+        hidden_states_norms = {}
+        for stage, num_channels in zip(self.out_features, self.channels):
+            hidden_states_norms[stage] = ConvNextV2LayerNorm(num_channels, data_format="channels_first")
+        self.hidden_states_norms = nn.ModuleDict(hidden_states_norms)
+
+        # initialize weights and apply final processing
+        self.post_init()
+
+    @can_return_tuple
+    @filter_output_hidden_states
+    @auto_docstring
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BackboneOutput:
+        r"""
+        Examples:
+
+        ```python
+        >>> from transformers import AutoImageProcessor, AutoBackbone
+        >>> import torch
+        >>> from PIL import Image
+        >>> import httpx
+        >>> from io import BytesIO
+
+        >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+        >>> with httpx.stream("GET", url) as response:
+        ...     image = Image.open(BytesIO(response.read()))
+
+        >>> processor = AutoImageProcessor.from_pretrained("facebook/convnextv2-tiny-1k-224")
+        >>> model = AutoBackbone.from_pretrained("facebook/convnextv2-tiny-1k-224")
+
+        >>> inputs = processor(image, return_tensors="pt")
+        >>> outputs = model(**inputs)
+        ```"""
+        kwargs["output_hidden_states"] = True  # required to extract layers for the stages
+
+        embedding_output = self.embeddings(pixel_values)
+        encoder_outputs: BaseModelOutputWithNoAttention = self.encoder(embedding_output, **kwargs)
+        hidden_states = encoder_outputs.hidden_states
+
+        feature_maps = []
+        for stage, hidden_state in zip(self.stage_names, hidden_states):
+            if stage in self.out_features:
+                hidden_state = self.hidden_states_norms[stage](hidden_state)
+                feature_maps.append(hidden_state)
+
+        return BackboneOutput(feature_maps=tuple(feature_maps), hidden_states=hidden_states)
+
+
+__all__ = ["ConvNextV2ForImageClassification", "ConvNextV2Model", "ConvNextV2PreTrainedModel", "ConvNextV2Backbone"]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/levit/__init__.py

@@ -0,0 +1,30 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+    from .configuration_levit import *
+    from .feature_extraction_levit import *
+    from .image_processing_levit import *
+    from .image_processing_pil_levit import *
+    from .modeling_levit import *
+else:
+    import sys
+
+    _file = globals()["__file__"]
+    sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/levit/modeling_levit.py

@@ -0,0 +1,665 @@
+# Copyright 2022 Meta Platforms, Inc. and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PyTorch LeViT model."""
+
+import itertools
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+
+from ... import initialization as init
+from ...modeling_outputs import (
+    BaseModelOutputWithNoAttention,
+    BaseModelOutputWithPoolingAndNoAttention,
+    ImageClassifierOutputWithNoAttention,
+    ModelOutput,
+)
+from ...modeling_utils import PreTrainedModel
+from ...utils import auto_docstring, logging
+from .configuration_levit import LevitConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+@auto_docstring(
+    custom_intro="""
+    Output type of [`LevitForImageClassificationWithTeacher`].
+    """
+)
+@dataclass
+class LevitForImageClassificationWithTeacherOutput(ModelOutput):
+    r"""
+    logits (`torch.FloatTensor` of shape `(batch_size, config.num_labels)`):
+        Prediction scores as the average of the `cls_logits` and `distillation_logits`.
+    cls_logits (`torch.FloatTensor` of shape `(batch_size, config.num_labels)`):
+        Prediction scores of the classification head (i.e. the linear layer on top of the final hidden state of the
+        class token).
+    distillation_logits (`torch.FloatTensor` of shape `(batch_size, config.num_labels)`):
+        Prediction scores of the distillation head (i.e. the linear layer on top of the final hidden state of the
+        distillation token).
+    """
+
+    logits: torch.FloatTensor | None = None
+    cls_logits: torch.FloatTensor | None = None
+    distillation_logits: torch.FloatTensor | None = None
+    hidden_states: tuple[torch.FloatTensor] | None = None
+
+
+class LevitConvEmbeddings(nn.Module):
+    """
+    LeViT Conv Embeddings with Batch Norm, used in the initial patch embedding layer.
+    """
+
+    def __init__(
+        self, in_channels, out_channels, kernel_size, stride, padding, dilation=1, groups=1, bn_weight_init=1
+    ):
+        super().__init__()
+        self.convolution = nn.Conv2d(
+            in_channels, out_channels, kernel_size, stride, padding, dilation=dilation, groups=groups, bias=False
+        )
+        self.batch_norm = nn.BatchNorm2d(out_channels)
+
+    def forward(self, embeddings):
+        embeddings = self.convolution(embeddings)
+        embeddings = self.batch_norm(embeddings)
+        return embeddings
+
+
+class LevitPatchEmbeddings(nn.Module):
+    """
+    LeViT patch embeddings, for final embeddings to be passed to transformer blocks. It consists of multiple
+    `LevitConvEmbeddings`.
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        self.embedding_layer_1 = LevitConvEmbeddings(
+            config.num_channels, config.hidden_sizes[0] // 8, config.kernel_size, config.stride, config.padding
+        )
+        self.activation_layer_1 = nn.Hardswish()
+
+        self.embedding_layer_2 = LevitConvEmbeddings(
+            config.hidden_sizes[0] // 8, config.hidden_sizes[0] // 4, config.kernel_size, config.stride, config.padding
+        )
+        self.activation_layer_2 = nn.Hardswish()
+
+        self.embedding_layer_3 = LevitConvEmbeddings(
+            config.hidden_sizes[0] // 4, config.hidden_sizes[0] // 2, config.kernel_size, config.stride, config.padding
+        )
+        self.activation_layer_3 = nn.Hardswish()
+
+        self.embedding_layer_4 = LevitConvEmbeddings(
+            config.hidden_sizes[0] // 2, config.hidden_sizes[0], config.kernel_size, config.stride, config.padding
+        )
+        self.num_channels = config.num_channels
+
+    def forward(self, pixel_values):
+        num_channels = pixel_values.shape[1]
+        if num_channels != self.num_channels:
+            raise ValueError(
+                "Make sure that the channel dimension of the pixel values match with the one set in the configuration."
+            )
+        embeddings = self.embedding_layer_1(pixel_values)
+        embeddings = self.activation_layer_1(embeddings)
+        embeddings = self.embedding_layer_2(embeddings)
+        embeddings = self.activation_layer_2(embeddings)
+        embeddings = self.embedding_layer_3(embeddings)
+        embeddings = self.activation_layer_3(embeddings)
+        embeddings = self.embedding_layer_4(embeddings)
+        return embeddings.flatten(2).transpose(1, 2)
+
+
+class MLPLayerWithBN(nn.Module):
+    def __init__(self, input_dim, output_dim, bn_weight_init=1):
+        super().__init__()
+        self.linear = nn.Linear(in_features=input_dim, out_features=output_dim, bias=False)
+        self.batch_norm = nn.BatchNorm1d(output_dim)
+
+    def forward(self, hidden_state):
+        hidden_state = self.linear(hidden_state)
+        hidden_state = self.batch_norm(hidden_state.flatten(0, 1)).reshape_as(hidden_state)
+        return hidden_state
+
+
+class LevitSubsample(nn.Module):
+    def __init__(self, stride, resolution):
+        super().__init__()
+        self.stride = stride
+        self.resolution = resolution
+
+    def forward(self, hidden_state):
+        batch_size, _, channels = hidden_state.shape
+        hidden_state = hidden_state.view(batch_size, self.resolution, self.resolution, channels)[
+            :, :: self.stride, :: self.stride
+        ].reshape(batch_size, -1, channels)
+        return hidden_state
+
+
+class LevitAttention(nn.Module):
+    def __init__(self, hidden_sizes, key_dim, num_attention_heads, attention_ratio, resolution):
+        super().__init__()
+        self.num_attention_heads = num_attention_heads
+        self.scale = key_dim**-0.5
+        self.key_dim = key_dim
+        self.attention_ratio = attention_ratio
+        self.out_dim_keys_values = attention_ratio * key_dim * num_attention_heads + key_dim * num_attention_heads * 2
+        self.out_dim_projection = attention_ratio * key_dim * num_attention_heads
+
+        self.queries_keys_values = MLPLayerWithBN(hidden_sizes, self.out_dim_keys_values)
+        self.activation = nn.Hardswish()
+        self.projection = MLPLayerWithBN(self.out_dim_projection, hidden_sizes, bn_weight_init=0)
+
+        points = list(itertools.product(range(resolution), range(resolution)))
+        len_points = len(points)
+        self.len_points = len_points
+        attention_offsets, indices = {}, []
+        for p1 in points:
+            for p2 in points:
+                offset = (abs(p1[0] - p2[0]), abs(p1[1] - p2[1]))
+                if offset not in attention_offsets:
+                    attention_offsets[offset] = len(attention_offsets)
+                indices.append(attention_offsets[offset])
+        self.indices = indices
+
+        self.attention_bias_cache = {}
+        self.attention_biases = torch.nn.Parameter(torch.zeros(num_attention_heads, len(attention_offsets)))
+        self.register_buffer(
+            "attention_bias_idxs", torch.LongTensor(indices).view(len_points, len_points), persistent=False
+        )
+
+    @torch.no_grad()
+    def train(self, mode=True):
+        super().train(mode)
+        if mode and self.attention_bias_cache:
+            self.attention_bias_cache = {}  # clear ab cache
+
+    def get_attention_biases(self, device):
+        if self.training:
+            return self.attention_biases[:, self.attention_bias_idxs]
+        else:
+            device_key = str(device)
+            if device_key not in self.attention_bias_cache:
+                self.attention_bias_cache[device_key] = self.attention_biases[:, self.attention_bias_idxs]
+            return self.attention_bias_cache[device_key]
+
+    def forward(self, hidden_state):
+        batch_size, seq_length, _ = hidden_state.shape
+        queries_keys_values = self.queries_keys_values(hidden_state)
+        query, key, value = queries_keys_values.view(batch_size, seq_length, self.num_attention_heads, -1).split(
+            [self.key_dim, self.key_dim, self.attention_ratio * self.key_dim], dim=3
+        )
+        query = query.permute(0, 2, 1, 3)
+        key = key.permute(0, 2, 1, 3)
+        value = value.permute(0, 2, 1, 3)
+
+        attention = query @ key.transpose(-2, -1) * self.scale + self.get_attention_biases(hidden_state.device)
+        attention = attention.softmax(dim=-1)
+        hidden_state = (attention @ value).transpose(1, 2).reshape(batch_size, seq_length, self.out_dim_projection)
+        hidden_state = self.projection(self.activation(hidden_state))
+        return hidden_state
+
+
+class LevitAttentionSubsample(nn.Module):
+    def __init__(
+        self,
+        input_dim,
+        output_dim,
+        key_dim,
+        num_attention_heads,
+        attention_ratio,
+        stride,
+        resolution_in,
+        resolution_out,
+    ):
+        super().__init__()
+        self.num_attention_heads = num_attention_heads
+        self.scale = key_dim**-0.5
+        self.key_dim = key_dim
+        self.attention_ratio = attention_ratio
+        self.out_dim_keys_values = attention_ratio * key_dim * num_attention_heads + key_dim * num_attention_heads
+        self.out_dim_projection = attention_ratio * key_dim * num_attention_heads
+        self.resolution_out = resolution_out
+        # resolution_in is the initial resolution, resolution_out is final resolution after downsampling
+        self.keys_values = MLPLayerWithBN(input_dim, self.out_dim_keys_values)
+        self.queries_subsample = LevitSubsample(stride, resolution_in)
+        self.queries = MLPLayerWithBN(input_dim, key_dim * num_attention_heads)
+        self.activation = nn.Hardswish()
+        self.projection = MLPLayerWithBN(self.out_dim_projection, output_dim)
+
+        self.attention_bias_cache = {}
+
+        points = list(itertools.product(range(resolution_in), range(resolution_in)))
+        points_ = list(itertools.product(range(resolution_out), range(resolution_out)))
+        len_points, len_points_ = len(points), len(points_)
+        self.len_points_ = len_points_
+        self.len_points = len_points
+        attention_offsets, indices = {}, []
+        for p1 in points_:
+            for p2 in points:
+                size = 1
+                offset = (abs(p1[0] * stride - p2[0] + (size - 1) / 2), abs(p1[1] * stride - p2[1] + (size - 1) / 2))
+                if offset not in attention_offsets:
+                    attention_offsets[offset] = len(attention_offsets)
+                indices.append(attention_offsets[offset])
+        self.indices = indices
+
+        self.attention_biases = torch.nn.Parameter(torch.zeros(num_attention_heads, len(attention_offsets)))
+        self.register_buffer(
+            "attention_bias_idxs", torch.LongTensor(indices).view(len_points_, len_points), persistent=False
+        )
+
+    @torch.no_grad()
+    def train(self, mode=True):
+        super().train(mode)
+        if mode and self.attention_bias_cache:
+            self.attention_bias_cache = {}  # clear ab cache
+
+    def get_attention_biases(self, device):
+        if self.training:
+            return self.attention_biases[:, self.attention_bias_idxs]
+        else:
+            device_key = str(device)
+            if device_key not in self.attention_bias_cache:
+                self.attention_bias_cache[device_key] = self.attention_biases[:, self.attention_bias_idxs]
+            return self.attention_bias_cache[device_key]
+
+    def forward(self, hidden_state):
+        batch_size, seq_length, _ = hidden_state.shape
+        key, value = (
+            self.keys_values(hidden_state)
+            .view(batch_size, seq_length, self.num_attention_heads, -1)
+            .split([self.key_dim, self.attention_ratio * self.key_dim], dim=3)
+        )
+        key = key.permute(0, 2, 1, 3)
+        value = value.permute(0, 2, 1, 3)
+
+        query = self.queries(self.queries_subsample(hidden_state))
+        query = query.view(batch_size, self.resolution_out**2, self.num_attention_heads, self.key_dim).permute(
+            0, 2, 1, 3
+        )
+
+        attention = query @ key.transpose(-2, -1) * self.scale + self.get_attention_biases(hidden_state.device)
+        attention = attention.softmax(dim=-1)
+        hidden_state = (attention @ value).transpose(1, 2).reshape(batch_size, -1, self.out_dim_projection)
+        hidden_state = self.projection(self.activation(hidden_state))
+        return hidden_state
+
+
+class LevitMLPLayer(nn.Module):
+    """
+    MLP Layer with `2X` expansion in contrast to ViT with `4X`.
+    """
+
+    def __init__(self, input_dim, hidden_dim):
+        super().__init__()
+        self.linear_up = MLPLayerWithBN(input_dim, hidden_dim)
+        self.activation = nn.Hardswish()
+        self.linear_down = MLPLayerWithBN(hidden_dim, input_dim)
+
+    def forward(self, hidden_state):
+        hidden_state = self.linear_up(hidden_state)
+        hidden_state = self.activation(hidden_state)
+        hidden_state = self.linear_down(hidden_state)
+        return hidden_state
+
+
+class LevitResidualLayer(nn.Module):
+    """
+    Residual Block for LeViT
+    """
+
+    def __init__(self, module, drop_rate):
+        super().__init__()
+        self.module = module
+        self.drop_rate = drop_rate
+
+    def forward(self, hidden_state):
+        if self.training and self.drop_rate > 0:
+            rnd = torch.rand(hidden_state.size(0), 1, 1, device=hidden_state.device)
+            rnd = rnd.ge_(self.drop_rate).div(1 - self.drop_rate).detach()
+            hidden_state = hidden_state + self.module(hidden_state) * rnd
+            return hidden_state
+        else:
+            hidden_state = hidden_state + self.module(hidden_state)
+            return hidden_state
+
+
+class LevitStage(nn.Module):
+    """
+    LeViT Stage consisting of `LevitMLPLayer` and `LevitAttention` layers.
+    """
+
+    def __init__(
+        self,
+        config,
+        idx,
+        hidden_sizes,
+        key_dim,
+        depths,
+        num_attention_heads,
+        attention_ratio,
+        mlp_ratio,
+        down_ops,
+        resolution_in,
+    ):
+        super().__init__()
+        self.layers = []
+        self.config = config
+        self.resolution_in = resolution_in
+        # resolution_in is the initial resolution, resolution_out is final resolution after downsampling
+        for _ in range(depths):
+            self.layers.append(
+                LevitResidualLayer(
+                    LevitAttention(hidden_sizes, key_dim, num_attention_heads, attention_ratio, resolution_in),
+                    self.config.drop_path_rate,
+                )
+            )
+            if mlp_ratio > 0:
+                hidden_dim = hidden_sizes * mlp_ratio
+                self.layers.append(
+                    LevitResidualLayer(LevitMLPLayer(hidden_sizes, hidden_dim), self.config.drop_path_rate)
+                )
+
+        if down_ops[0] == "Subsample":
+            self.resolution_out = (self.resolution_in - 1) // down_ops[5] + 1
+            self.layers.append(
+                LevitAttentionSubsample(
+                    *self.config.hidden_sizes[idx : idx + 2],
+                    key_dim=down_ops[1],
+                    num_attention_heads=down_ops[2],
+                    attention_ratio=down_ops[3],
+                    stride=down_ops[5],
+                    resolution_in=resolution_in,
+                    resolution_out=self.resolution_out,
+                )
+            )
+            self.resolution_in = self.resolution_out
+            if down_ops[4] > 0:
+                hidden_dim = self.config.hidden_sizes[idx + 1] * down_ops[4]
+                self.layers.append(
+                    LevitResidualLayer(
+                        LevitMLPLayer(self.config.hidden_sizes[idx + 1], hidden_dim), self.config.drop_path_rate
+                    )
+                )
+
+        self.layers = nn.ModuleList(self.layers)
+
+    def get_resolution(self):
+        return self.resolution_in
+
+    def forward(self, hidden_state):
+        for layer in self.layers:
+            hidden_state = layer(hidden_state)
+        return hidden_state
+
+
+class LevitEncoder(nn.Module):
+    """
+    LeViT Encoder consisting of multiple `LevitStage` stages.
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        resolution = self.config.image_size // self.config.patch_size
+        self.stages = []
+        self.config.down_ops.append([""])
+
+        for stage_idx in range(len(config.depths)):
+            stage = LevitStage(
+                config,
+                stage_idx,
+                config.hidden_sizes[stage_idx],
+                config.key_dim[stage_idx],
+                config.depths[stage_idx],
+                config.num_attention_heads[stage_idx],
+                config.attention_ratio[stage_idx],
+                config.mlp_ratio[stage_idx],
+                config.down_ops[stage_idx],
+                resolution,
+            )
+            resolution = stage.get_resolution()
+            self.stages.append(stage)
+
+        self.stages = nn.ModuleList(self.stages)
+
+    def forward(self, hidden_state, output_hidden_states=False, return_dict=True):
+        all_hidden_states = () if output_hidden_states else None
+
+        for stage in self.stages:
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_state,)
+            hidden_state = stage(hidden_state)
+
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_state,)
+        if not return_dict:
+            return tuple(v for v in [hidden_state, all_hidden_states] if v is not None)
+
+        return BaseModelOutputWithNoAttention(last_hidden_state=hidden_state, hidden_states=all_hidden_states)
+
+
+class LevitClassificationLayer(nn.Module):
+    """
+    LeViT Classification Layer
+    """
+
+    def __init__(self, input_dim, output_dim):
+        super().__init__()
+        self.batch_norm = nn.BatchNorm1d(input_dim)
+        self.linear = nn.Linear(input_dim, output_dim)
+
+    def forward(self, hidden_state):
+        hidden_state = self.batch_norm(hidden_state)
+        logits = self.linear(hidden_state)
+        return logits
+
+
+@auto_docstring
+class LevitPreTrainedModel(PreTrainedModel):
+    config: LevitConfig
+    base_model_prefix = "levit"
+    main_input_name = "pixel_values"
+    input_modalities = ("image",)
+    _no_split_modules = ["LevitResidualLayer"]
+
+    def _init_weights(self, module):
+        super()._init_weights(module)
+        if isinstance(module, LevitAttention):
+            init.copy_(
+                module.attention_bias_idxs, torch.LongTensor(module.indices).view(module.len_points, module.len_points)
+            )
+        elif isinstance(module, LevitAttentionSubsample):
+            init.copy_(
+                module.attention_bias_idxs,
+                torch.LongTensor(module.indices).view(module.len_points_, module.len_points),
+            )
+
+
+@auto_docstring
+class LevitModel(LevitPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+        self.patch_embeddings = LevitPatchEmbeddings(config)
+        self.encoder = LevitEncoder(config)
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @auto_docstring
+    def forward(
+        self,
+        pixel_values: torch.FloatTensor | None = None,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+        **kwargs,
+    ) -> tuple | BaseModelOutputWithPoolingAndNoAttention:
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+        if pixel_values is None:
+            raise ValueError("You have to specify pixel_values")
+
+        embeddings = self.patch_embeddings(pixel_values)
+        encoder_outputs = self.encoder(
+            embeddings,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        last_hidden_state = encoder_outputs[0]
+
+        # global average pooling, (batch_size, seq_length, hidden_sizes) -> (batch_size, hidden_sizes)
+        pooled_output = last_hidden_state.mean(dim=1)
+
+        if not return_dict:
+            return (last_hidden_state, pooled_output) + encoder_outputs[1:]
+
+        return BaseModelOutputWithPoolingAndNoAttention(
+            last_hidden_state=last_hidden_state,
+            pooler_output=pooled_output,
+            hidden_states=encoder_outputs.hidden_states,
+        )
+
+
+@auto_docstring(
+    custom_intro="""
+    Levit Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for
+    ImageNet.
+    """
+)
+class LevitForImageClassification(LevitPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+        self.num_labels = config.num_labels
+        self.levit = LevitModel(config)
+
+        # Classifier head
+        self.classifier = (
+            LevitClassificationLayer(config.hidden_sizes[-1], config.num_labels)
+            if config.num_labels > 0
+            else torch.nn.Identity()
+        )
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @auto_docstring
+    def forward(
+        self,
+        pixel_values: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+        **kwargs,
+    ) -> tuple | ImageClassifierOutputWithNoAttention:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Labels for computing the image classification/regression loss. Indices should be in `[0, ...,
+            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+        """
+        return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+        outputs = self.levit(pixel_values, output_hidden_states=output_hidden_states, return_dict=return_dict)
+
+        sequence_output = outputs[0]
+        sequence_output = sequence_output.mean(1)
+        logits = self.classifier(sequence_output)
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(labels, logits, self.config)
+
+        if not return_dict:
+            output = (logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return ImageClassifierOutputWithNoAttention(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+        )
+
+
+@auto_docstring(
+    custom_intro="""
+    LeViT Model transformer with image classification heads on top (a linear layer on top of the final hidden state and
+    a linear layer on top of the final hidden state of the distillation token) e.g. for ImageNet. .. warning::
+           This model supports inference-only. Fine-tuning with distillation (i.e. with a teacher) is not yet
+           supported.
+    """
+)
+class LevitForImageClassificationWithTeacher(LevitPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+        self.num_labels = config.num_labels
+        self.levit = LevitModel(config)
+
+        # Classifier head
+        self.classifier = (
+            LevitClassificationLayer(config.hidden_sizes[-1], config.num_labels)
+            if config.num_labels > 0
+            else torch.nn.Identity()
+        )
+        self.classifier_distill = (
+            LevitClassificationLayer(config.hidden_sizes[-1], config.num_labels)
+            if config.num_labels > 0
+            else torch.nn.Identity()
+        )
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @auto_docstring
+    def forward(
+        self,
+        pixel_values: torch.FloatTensor | None = None,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+        **kwargs,
+    ) -> tuple | LevitForImageClassificationWithTeacherOutput:
+        return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+        outputs = self.levit(pixel_values, output_hidden_states=output_hidden_states, return_dict=return_dict)
+
+        sequence_output = outputs[0]
+        sequence_output = sequence_output.mean(1)
+        cls_logits, distill_logits = self.classifier(sequence_output), self.classifier_distill(sequence_output)
+        logits = (cls_logits + distill_logits) / 2
+
+        if not return_dict:
+            output = (logits, cls_logits, distill_logits) + outputs[2:]
+            return output
+
+        return LevitForImageClassificationWithTeacherOutput(
+            logits=logits,
+            cls_logits=cls_logits,
+            distillation_logits=distill_logits,
+            hidden_states=outputs.hidden_states,
+        )
+
+
+__all__ = [
+    "LevitForImageClassification",
+    "LevitForImageClassificationWithTeacher",
+    "LevitModel",
+    "LevitPreTrainedModel",
+]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/megatron_gpt2/checkpoint_reshaping_and_interoperability.py

@@ -0,0 +1,925 @@
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import importlib.util
+import json
+import os
+import re
+import sys
+import types
+
+import torch
+from huggingface_hub import split_torch_state_dict_into_shards
+from packaging import version
+
+from transformers import AutoTokenizer, GPT2Config
+from transformers.modeling_utils import WEIGHTS_INDEX_NAME, WEIGHTS_NAME
+from transformers.utils import check_torch_load_is_safe
+
+
+def add_checkpointing_args(parser):
+    parser.add_argument("--megatron-path", type=str, default=None, help="Base directory of Megatron repository")
+    parser.add_argument(
+        "--convert_checkpoint_from_megatron_to_transformers",
+        action="store_true",
+        help=(
+            "If True, convert a Megatron checkpoint to a Transformers checkpoint. "
+            "If False, convert a Transformers checkpoint to a Megatron checkpoint."
+        ),
+    )
+    parser.add_argument(
+        "--load_path",
+        type=str,
+        required=True,
+        help="Path to the checkpoint to convert.",
+    )
+    parser.add_argument(
+        "--save_path",
+        type=str,
+        required=True,
+        help="Path to the converted checkpoint.",
+    )
+    parser.add_argument("--print-checkpoint-structure", action="store_true")
+    return parser
+
+
+def add_megatron_checkpoint_args(parser):
+    parser.add_argument(
+        "--target_tensor_model_parallel_size",
+        type=int,
+        default=1,
+        help=(
+            "The tensor model parallel size of the converted checkpoint. "
+            "Only used when converting a Transformers checkpoint to a Megatron checkpoint."
+        ),
+    )
+    parser.add_argument(
+        "--target_pipeline_model_parallel_size",
+        type=int,
+        default=1,
+        help=(
+            "The pipeline model parallel size of the converted checkpoint. "
+            "Only used when converting a Transformers checkpoint to a Megatron checkpoint."
+        ),
+    )
+    parser.add_argument(
+        "--target_data_parallel_size",
+        type=int,
+        default=1,
+        help=(
+            "The data parallel size of the converted checkpoint. "
+            "Only used when converting a Transformers checkpoint to a Megatron checkpoint."
+        ),
+    )
+    parser.add_argument(
+        "--target_params_dtype",
+        type=str,
+        default="fp32",
+        help=(
+            "The dtype of the converted checkpoint. "
+            "Only used when converting a Transformers checkpoint to a Megatron checkpoint."
+        ),
+    )
+    parser.add_argument(
+        "--make_vocab_size_divisible_by",
+        type=int,
+        default=128,
+        help=(
+            "Pad the vocab size to be divisible by this value. "
+            "This is added for computational efficiency reasons. "
+            "Only used when converting a Transformers checkpoint to a Megatron checkpoint."
+        ),
+    )
+    parser.add_argument(
+        "--use_distributed_optimizer",
+        action="store_true",
+        help=(
+            "If True, use the distributed optimizer. "
+            "Only used when converting a Transformers checkpoint to a Megatron checkpoint."
+        ),
+    )
+    return parser
+
+
+def add_transformers_checkpoint_args(parser):
+    parser.add_argument(
+        "--tokenizer_name",
+        type=str,
+        default=None,
+        help=(
+            "The name of the pre-trained tokenizer to save. "
+            "If not None, the tokenizer will be saved. "
+            "Only used when converting a Megatron checkpoint to a Transformers checkpoint."
+        ),
+    )
+    parser.add_argument(
+        "--max_shard_size",
+        type=str,
+        default="10GB",
+        help=(
+            "The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size "
+            "lower than this size. If expressed as a string, needs to be digits followed by a unit (like `5MB`). "
+            "Only used when converting a Megatron checkpoint to a Transformers checkpoint."
+        ),
+    )
+
+    return parser
+
+
+# The simple map of names for "automated" rules.
+megatron_to_transformers = {
+    "attention.dense": ".attn.c_proj.",
+    "self_attention.dense": ".attn.c_proj.",
+    "mlp.dense_h_to_4h": ".mlp.c_fc.",
+    "mlp.dense_4h_to_h": ".mlp.c_proj.",
+}
+transformers_to_megatron = {v[1:-1]: k for k, v in megatron_to_transformers.items()}
+
+tensor_parallel_params = [
+    # megatron-lm layers to merge across tp ranks
+    "self_attention.query_key_value.weight",
+    "self_attention.query_key_value.bias",
+    "self_attention.dense.weight",
+    "mlp.dense_h_to_4h.weight",
+    "mlp.dense_h_to_4h.bias",
+    "mlp.dense_4h_to_h.weight",
+    # deprecated
+    "attention.query_key_value.weight",
+    "attention.query_key_value.bias",
+    "attention.dense.weight",
+    # transformers layers to split across tp ranks
+    "attn.c_attn.weight",
+    "attn.c_attn.bias",
+    "attn.c_proj.weight",
+    "mlp.c_fc.weight",
+    "mlp.c_fc.bias",
+    "mlp.c_proj.weight",
+]
+
+
+def recursive_print(name, val, spaces=0):
+    """
+    Recursively print the structure of a checkpoint. This function is taken from `convert_megatron_gpt2_checkpoint.py`
+
+    Args:
+        name (str): the name of the current tensor parameter
+        val (Tuple(int)): the shape of the current tensor parameter
+        spaces (int): the number of spaces to print before the output for a nested structure
+    """
+    # Format the message.
+    if name is None:
+        msg = None
+    else:
+        fmt = "." * max(0, spaces - 2) + "# {:" + str(50 - spaces) + "s}"
+        msg = fmt.format(name)
+
+    # Print and recurse (if needed).
+    if isinstance(val, dict):
+        if msg is not None:
+            print(msg)
+        for k in val:
+            recursive_print(k, val[k], spaces + 2)
+    elif isinstance(val, torch.Tensor):
+        print(msg, ":", val.size())
+    else:
+        print(msg, ":", val)
+
+
+def megatron_to_transformers_fix_query_key_value_ordering(
+    param, checkpoint_version, num_splits, num_heads, hidden_size
+):
+    """
+    Permutes layout of param tensor to [num_splits * num_heads * hidden_size, :] for compatibility with later versions
+    of NVIDIA Megatron-LM. The inverse operation is performed inside Megatron-LM to read checkpoints:
+    https://github.com/NVIDIA/Megatron-LM/blob/v2.4/megatron/checkpointing.py#L209 If param is the weight tensor of the
+    self-attention block, the returned tensor will have to be transposed one more time to be read by HuggingFace GPT2.
+    This function is taken from `convert_megatron_gpt2_checkpoint.py`
+
+    Args:
+        param (torch.Tensor): the tensor to permute
+        checkpoint_version (int): the version of the checkpoint.
+        num_splits (int): the number of projections, usually 3 for (Query, Key, Value)
+        num_heads (int): the number of attention heads
+        hidden_size (int): the hidden size per head
+    """
+
+    input_shape = param.size()
+    if checkpoint_version == 1.0:
+        # version 1.0 stores [num_heads * hidden_size * num_splits, :]
+        saved_shape = (num_heads, hidden_size, num_splits) + input_shape[1:]
+        param = param.view(*saved_shape)
+        param = param.transpose(0, 2)
+        param = param.transpose(1, 2).contiguous()
+    elif checkpoint_version >= 2.0:
+        # other versions store [num_heads * num_splits * hidden_size, :]
+        saved_shape = (num_heads, num_splits, hidden_size) + input_shape[1:]
+        param = param.view(*saved_shape)
+        param = param.transpose(0, 1).contiguous()
+    param = param.view(*input_shape)
+    return param
+
+
+def transformers_to_megatron_fix_query_key_value_ordering(
+    param, checkpoint_version, num_splits, num_heads, hidden_size
+):
+    """
+    Permutes layout of param tensor to the one compatible with respective NVIDIA Megatron-LM checkpoint versions. Input
+    is [num_splits * num_heads * hidden_size, :] and output is [num_heads * hidden_size * num_splits, :] for version
+    1.0 and [num_heads * num_splits * hidden_size, :] for version 2.0 and later. If param is the weight tensor of the
+    self-attention block, the param needs to be already transposed before calling this function.
+
+    Args:
+        param (torch.Tensor): the tensor to permute
+        checkpoint_version (int): the version of the checkpoint.
+        num_splits (int): the number of projections, usually 3 for (Query, Key, Value)
+        num_heads (int): the number of attention heads
+        hidden_size (int): the hidden size per head
+    """
+
+    # Input is [num_splits * num_heads * hidden_size, :]
+    input_shape = param.size()
+    if checkpoint_version == 1.0:
+        # version 1.0 stores [num_heads * hidden_size * num_splits, :]
+        current_shape = (num_splits, num_heads, hidden_size) + input_shape[1:]
+        param = param.view(*current_shape)
+        param = param.transpose(0, 2)
+        param = param.transpose(1, 2).contiguous()
+    elif checkpoint_version >= 2.0:
+        # other versions store [num_heads * num_splits * hidden_size, :]
+        current_shape = (num_splits, num_heads, hidden_size) + input_shape[1:]
+        param = param.view(*current_shape)
+        param = param.transpose(0, 1).contiguous()
+    param = param.view(*input_shape)
+    return param
+
+
+def merge_transformers_sharded_states(path, num_checkpoints):
+    """
+    Merge sharded checkpoints from transformers into a single checkpoint.
+
+    Args:
+        path (str): the path to the sharded checkpoints
+        num_checkpoints (int): the number of checkpoints to merge
+    """
+    state_dict = {}
+    for i in range(1, num_checkpoints + 1):
+        checkpoint_path = os.path.join(path, f"pytorch_model-{i:05d}-of-{num_checkpoints:05d}.bin")
+        check_torch_load_is_safe()
+        current_chunk = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
+        state_dict.update(current_chunk)
+    return state_dict
+
+
+def get_megatron_sharded_states(args, tp_size, pp_size, pp_rank):
+    """
+    Get sharded checkpoints from NVIDIA Megatron-LM checkpoint based on the provided tensor parallel size, pipeline
+    parallel size and pipeline parallel rank.
+
+    Args:
+        args (argparse.Namespace): the arguments to the script
+        tp_size (int): the tensor parallel size
+        pp_size (int): the pipeline parallel size
+        pp_rank (int): the pipeline parallel rank
+    """
+    tp_state_dicts = []
+    for i in range(tp_size):
+        sub_dir_name = f"mp_rank_{i:02d}" if pp_size == 1 else f"mp_rank_{i:02d}_{pp_rank:03d}"
+        for checkpoint_name in ["model_optim_rng.pt", "model_rng.pt"]:
+            checkpoint_path = os.path.join(args.load_path, sub_dir_name, checkpoint_name)
+            if os.path.isfile(checkpoint_path):
+                break
+        check_torch_load_is_safe()
+        state_dict = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
+        tp_state_dicts.append(state_dict)
+    return tp_state_dicts
+
+
+def get_element_from_dict_by_path(d, path):
+    """
+    Get element from dictionary by path. If element is not present, recursively add empty dictionaries.
+
+    Args:
+        d (dict): the dictionary to get the element from
+        path (list): the path to the element which is delimited by "."
+    """
+    path = path.split(".")
+    for k in path:
+        if k not in d:
+            d[k] = {}
+        d = d[k]
+    return d
+
+
+def convert_checkpoint_from_megatron_to_transformers(args):
+    """
+    Convert NVIDIA Megatron-LM checkpoint to HuggingFace Transformers checkpoint. This handles Megatron checkpoints
+    with different tensor parallelism and pipeline parallelism sizes. It saves the converted checkpoint into shards
+    using HuggingFace Transformers checkpoint sharding functionality. This greatly extends the functionality of
+    `convert_megatron_gpt2_checkpoint.py`
+
+    Args:
+        args (argparse.Namespace): the arguments to the script
+    """
+    # Load Megatron-LM checkpoint arguments from the state dict
+    sub_dirs = os.listdir(args.load_path)
+    possible_sub_dirs = ["mp_rank_00", "mp_rank_00_000"]
+    for sub_dir in possible_sub_dirs:
+        if sub_dir in sub_dirs:
+            rank0_checkpoint_name = os.listdir(os.path.join(args.load_path, sub_dir))[0]
+            rank0_checkpoint_path = os.path.join(args.load_path, sub_dir, rank0_checkpoint_name)
+            break
+    print(f"Loading Megatron-LM checkpoint arguments from: {rank0_checkpoint_path}")
+    check_torch_load_is_safe()
+    state_dict = torch.load(rank0_checkpoint_path, map_location="cpu", weights_only=True)
+    megatron_args = state_dict.get("args", None)
+    if megatron_args is None:
+        raise ValueError(
+            "Megatron-LM checkpoint does not contain arguments. This utility only supports Megatron-LM checkpoints"
+            " containing all the megatron arguments. This is because it loads all config related to model"
+            " architecture, the tensor and pipeline model parallel size from the checkpoint instead of user having to"
+            " manually specify all the details. Please save Megatron-LM checkpoint along with all the megatron"
+            " arguments to use this utility."
+        )
+
+    # Create Transformers GPT2 config from Megatron-LM arguments
+    if megatron_args is not None:
+        if megatron_args.bias_gelu_fusion:
+            activation_function = "gelu_fast"
+        elif megatron_args.openai_gelu:
+            activation_function = "gelu_new"
+        else:
+            activation_function = "gelu"
+    else:
+        # in the very early days this used to be "gelu_new"
+        activation_function = "gelu_new"
+    vocab_size = (
+        megatron_args.padded_vocab_size
+        if getattr(megatron_args, "orig_vocab_size", None) is None
+        else megatron_args.orig_vocab_size
+    )
+    print(vocab_size)
+
+    config = GPT2Config(
+        vocab_size=vocab_size,
+        n_positions=megatron_args.max_position_embeddings,
+        n_embd=megatron_args.hidden_size,
+        n_layer=megatron_args.num_layers,
+        n_head=megatron_args.num_attention_heads,
+        n_inner=megatron_args.ffn_hidden_size,
+        activation_function=activation_function,
+        resid_pdrop=0.1,
+        embd_pdrop=0.1,
+        attn_pdrop=0.1,
+        layer_norm_epsilon=1e-5,
+        initializer_range=0.02,
+        summary_type="cls_index",
+        summary_use_proj=True,
+        summary_activation=None,
+        summary_proj_to_labels=True,
+        summary_first_dropout=0.1,
+        scale_attn_weights=True,
+        use_cache=True,
+        bos_token_id=vocab_size - 1,
+        eos_token_id=vocab_size - 1,
+        architectures=["GPT2LMHeadModel"],
+    )
+
+    output_state_dict = {}
+
+    checkpoint_version = state_dict.get("checkpoint_version", 0.0)
+    tp_size = megatron_args.tensor_model_parallel_size
+    pp_size = megatron_args.pipeline_model_parallel_size
+    dtype = torch.float32
+    # The regex to extract layer names.
+    layer_re = re.compile(r"layers\.(\d+)\.([a-z0-9_.]+)\.([a-z]+)")
+
+    # Convert.
+    print("Converting")
+
+    # Embeddings
+    print("Converting embeddings")
+    tp_state_dicts = get_megatron_sharded_states(args, tp_size, pp_size, 0)
+
+    # Convert and store the position embeddings.
+    position_embeddings = get_element_from_dict_by_path(
+        tp_state_dicts[0], "model.language_model.embedding.position_embeddings.weight"
+    )
+    output_state_dict["transformer.wpe.weight"] = position_embeddings.to(dtype)
+
+    # Convert and store the word embeddings.
+    word_embeddings = torch.cat(
+        [
+            get_element_from_dict_by_path(
+                tp_state_dicts[tp_rank], "model.language_model.embedding.word_embeddings.weight"
+            )
+            for tp_rank in range(tp_size)
+        ],
+        dim=0,
+    )
+    word_embeddings = word_embeddings[:vocab_size].to(dtype)
+    output_state_dict["transformer.wte.weight"] = word_embeddings
+
+    # Transformer Layers
+    print("Converting transformer layers")
+    # The number of heads.
+    heads = config.n_head
+    # The hidden_size per head.
+    hidden_size_per_head = config.n_embd // config.n_head
+    n_positions = config.n_positions
+    num_layers = config.num_hidden_layers // pp_size
+
+    for pp_rank in range(pp_size):
+        if pp_size > 0:
+            print(f"Converting pipeline parallel rank {pp_rank}")
+            tp_state_dicts = get_megatron_sharded_states(args, tp_size, pp_size, pp_rank)
+
+        # The transformer.
+        path = (
+            "model.language_model.transformer"
+            if "transformer" in get_element_from_dict_by_path(tp_state_dicts[0], "model.language_model")
+            else "model.language_model.encoder"
+        )
+        # Extract the layers.
+        for key, val in get_element_from_dict_by_path(tp_state_dicts[0], path).items():
+            # Match the name.
+            m = layer_re.match(key)
+            # Stop if that's not a layer
+            if m is None:
+                break
+
+            # The index of the layer.
+            layer_idx = int(m.group(1)) + pp_rank * num_layers
+            # The name of the operation.
+            op_name = m.group(2)
+            # Is it a weight or a bias?
+            weight_or_bias = m.group(3)
+
+            # The name of the layer.
+            layer_name = f"transformer.h.{layer_idx}"
+
+            if op_name + "." + weight_or_bias not in tensor_parallel_params:
+                params = val.to(dtype)
+            else:
+                dim = 1 if op_name in ["self_attention.dense", "mlp.dense_4h_to_h", "attention.dense"] else 0
+                params = torch.cat(
+                    [val]
+                    + [
+                        get_element_from_dict_by_path(tp_state_dicts[tp_rank], f"{path}")[key]
+                        for tp_rank in range(1, tp_size)
+                    ],
+                    dim=dim,
+                ).to(dtype)
+
+            # For layernorm(s), simply store the layer norm.
+            if op_name.endswith("layernorm"):
+                ln_name = "ln_1" if op_name.startswith("input") else "ln_2"
+                output_state_dict[layer_name + "." + ln_name + "." + weight_or_bias] = params
+
+            # Transpose the QKV matrix.
+            elif (
+                op_name == "attention.query_key_value" or op_name == "self_attention.query_key_value"
+            ) and weight_or_bias == "weight":
+                # Insert a tensor of 1x1xDxD bias.
+                causal_mask = torch.tril(torch.ones((n_positions, n_positions), dtype=dtype)).view(
+                    1, 1, n_positions, n_positions
+                )
+                output_state_dict[layer_name + ".attn.bias"] = causal_mask
+
+                # Insert a "dummy" tensor for masked_bias.
+                masked_bias = torch.tensor(-1e4, dtype=dtype)
+                output_state_dict[layer_name + ".attn.masked_bias"] = masked_bias
+
+                out_val = megatron_to_transformers_fix_query_key_value_ordering(
+                    params,
+                    checkpoint_version,
+                    3,
+                    heads,
+                    hidden_size_per_head,
+                )
+                # Megatron stores (3*D) x D but transformers-GPT2 expects D x 3*D.
+                out_val = out_val.transpose(0, 1).contiguous()
+                # Store.
+                output_state_dict[layer_name + ".attn.c_attn.weight"] = out_val
+
+            # Transpose the bias.
+            elif (
+                op_name == "attention.query_key_value" or op_name == "self_attention.query_key_value"
+            ) and weight_or_bias == "bias":
+                out_val = megatron_to_transformers_fix_query_key_value_ordering(
+                    params, checkpoint_version, 3, heads, hidden_size_per_head
+                )
+                # Store. No change of shape.
+                output_state_dict[layer_name + ".attn.c_attn.bias"] = out_val
+
+            # Transpose the weights.
+            elif weight_or_bias == "weight":
+                out_name = megatron_to_transformers[op_name]
+                output_state_dict[layer_name + out_name + "weight"] = params.transpose(0, 1)
+
+            # Copy the bias.
+            elif weight_or_bias == "bias":
+                out_name = megatron_to_transformers[op_name]
+                output_state_dict[layer_name + out_name + "bias"] = params
+
+    if config.n_layer != (layer_idx + 1):
+        raise ValueError(f"Expected {config.n_layer} layers but found {layer_idx + 1}")
+
+    # The final layernorm.
+    print("Converting final layernorm")
+    params = get_element_from_dict_by_path(tp_state_dicts[0], str(path))
+    output_state_dict["transformer.ln_f.weight"] = params["final_layernorm.weight"].to(dtype)
+    output_state_dict["transformer.ln_f.bias"] = params["final_layernorm.bias"].to(dtype)
+
+    # For LM head, transformers' wants the matrix to weight embeddings.
+    print("Converting LM head")
+    output_state_dict["lm_head.weight"] = word_embeddings.to(dtype)
+
+    # It should be done!
+    print("Conversion from Megatron-LM to Transformers is done!")
+
+    # Print the structure of converted state dict.
+    if args.print_checkpoint_structure:
+        recursive_print(None, output_state_dict)
+
+    # Add tokenizer class info to config
+    # see https://github.com/huggingface/transformers/issues/13906)
+
+    if args.tokenizer_name is None:
+        tokenizer_name = "openai-community/gpt2"
+    else:
+        tokenizer_name = args.tokenizer_name
+
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
+    tokenizer_class = type(tokenizer).__name__
+    config.tokenizer_class = tokenizer_class
+
+    # Store the config to file.
+    print("Saving config")
+    config.save_pretrained(args.save_path)
+
+    # Save tokenizer based on args
+    if args.tokenizer_name is not None:
+        print(f"Adding {tokenizer_class} tokenizer files")
+        tokenizer.save_pretrained(args.save_path)
+
+    # Store the state_dict to file.
+    max_shard_size = int(args.max_shard_size) if args.max_shard_size.isdigit() else args.max_shard_size
+    state_dict_split = split_torch_state_dict_into_shards(output_state_dict, max_shard_size=max_shard_size)
+    shards = index = None
+    for tensors in state_dict_split.filename_to_tensors.values():
+        shards = {tensor: state_dict[tensor] for tensor in tensors}
+    if state_dict_split.is_sharded:
+        index = {
+            "metadata": state_dict_split.metadata,
+            "weight_map": state_dict_split.tensor_to_filename,
+        }
+
+    # Save the model
+    for shard_file, shard in shards.items():
+        torch.save(shard, os.path.join(args.save_path, shard_file))
+
+    if index is None:
+        print(f"Model weights saved in {os.path.join(args.save_path, WEIGHTS_NAME)}")
+    else:
+        save_index_file = os.path.join(args.save_path, WEIGHTS_INDEX_NAME)
+        # Save the index as well
+        with open(save_index_file, "w", encoding="utf-8") as f:
+            content = json.dumps(index, indent=2, sort_keys=True) + "\n"
+            f.write(content)
+        print(
+            f"The model is bigger than the maximum size per checkpoint ({args.max_shard_size}) and is going to be "
+            f"split in {len(shards)} checkpoint shards. You can find where each parameters has been saved in the "
+            f"index located at {save_index_file}."
+        )
+
+
+def convert_checkpoint_from_transformers_to_megatron(args):
+    """
+    Convert a checkpoint from HuggingFace Transformers to Megatron-LM. This allows converted checkpoints with variable
+    tensor parallelism and pipeline parallelism sizes. It takes as input a checkpoint from HuggingFace Transformers
+    which can have multiple shards.
+
+    Args:
+        args (argparse.Namespace): the arguments to the script
+
+    """
+    os.makedirs(args.save_path, exist_ok=True)
+    # Search in directory above this
+    sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir)))
+    if args.megatron_path is not None:
+        sys.path.insert(0, args.megatron_path)
+
+    megatron_exists = importlib.util.find_spec("megatron") is not None
+    if megatron_exists:
+        from megatron.core import package_info
+
+        if version.parse(package_info.__version__) >= version.parse("0.6.0"):
+            from megatron.training.tokenizer.tokenizer import _vocab_size_with_padding
+        else:
+            from megatron.tokenizer.tokenizer import _vocab_size_with_padding
+
+    else:
+        print("Unable to import Megatron, please specify the path to Megatron using --megatron-path. Exiting.")
+        exit(1)
+
+    # load the transformers model state dict and config
+    sub_dirs = [x for x in os.listdir(args.load_path) if x.startswith("pytorch_model")]
+    if len(sub_dirs) == 1:
+        checkpoint_name = "pytorch_model.bin"
+        check_torch_load_is_safe()
+        state_dict = torch.load(os.path.join(args.load_path, checkpoint_name), map_location="cpu", weights_only=True)
+    else:
+        num_checkpoints = len(sub_dirs) - 1
+        state_dict = merge_transformers_sharded_states(args.load_path, num_checkpoints)
+
+    config = GPT2Config.from_pretrained(args.load_path)
+
+    # Saving the tracker file
+    tracker_filepath = os.path.join(args.save_path, "latest_checkpointed_iteration.txt")
+    with open(tracker_filepath, "w") as f:
+        f.write("release")
+
+    # create `release` dir in args.load_path
+    release_dir = os.path.join(args.save_path, "release")
+    os.makedirs(release_dir, exist_ok=True)
+
+    # megatron args
+    megatron_args = {
+        "orig_vocab_size": config.vocab_size,
+        "max_position_embeddings": config.n_positions,
+        "hidden_size": config.n_embd,
+        "num_layers": config.n_layer,
+        "num_attention_heads": config.n_head,
+        "ffn_hidden_size": config.n_inner,
+        "tensor_model_parallel_size": args.target_tensor_model_parallel_size,
+        "pipeline_model_parallel_size": args.target_pipeline_model_parallel_size,
+        "data_parallel_size": args.target_data_parallel_size,
+        "make_vocab_size_divisible_by": args.make_vocab_size_divisible_by,
+        "rank": 0,
+        "tokenizer_type": "GPT2BPETokenizer",
+    }
+
+    if config.activation_function == "gelu":
+        megatron_args["bias_gelu_fusion"] = False
+        megatron_args["openai_gelu"] = False
+    elif config.activation_function == "gelu_fast":
+        megatron_args["bias_gelu_fusion"] = True
+        megatron_args["openai_gelu"] = False
+    elif config.activation_function == "gelu_new":
+        megatron_args["bias_gelu_fusion"] = False
+        megatron_args["openai_gelu"] = True
+
+    margs = types.SimpleNamespace()
+    for k, v in megatron_args.items():
+        setattr(margs, k, v)
+
+    # params dtype
+    if args.target_params_dtype == "fp16":
+        dtype = torch.float16
+    elif args.target_params_dtype == "bf16":
+        dtype = torch.bfloat16
+    else:
+        dtype = torch.float32
+    setattr(margs, "params_dtype", dtype)
+
+    # save dummy optim state dict
+    dummy_optim_state_dict = {}
+    dummy_optim_state_dict["optimizer"] = {
+        "step": 0,
+        "param_groups": [
+            {
+                "lr": 0.0,
+                "beta1": 0.0,
+                "beta2": 0.0,
+                "eps": 0.0,
+                "weight_decay": 0.0,
+                "correct_bias": False,
+                "params": [],
+            }
+        ],
+    }
+    if args.use_distributed_optimizer:
+        for i in range(args.target_pipeline_model_parallel_size):
+            for j in range(args.target_tensor_model_parallel_size):
+                for k in range(args.target_data_parallel_size):
+                    if args.target_pipeline_model_parallel_size == 1:
+                        checkpoint_dir = f"mp_rank_{j:02d}_{k:03d}"
+                    else:
+                        checkpoint_dir = f"mp_rank_{j:02d}_{i:03d}_{k:03d}"
+                    checkpoint_dir = os.path.join(release_dir, checkpoint_dir)
+                    os.makedirs(checkpoint_dir, exist_ok=True)
+                    torch.save(
+                        dummy_optim_state_dict,
+                        os.path.join(checkpoint_dir, "optim.pt"),
+                    )
+
+    # Convert.
+    print("Converting")
+    output_state_dict = []
+    for i in range(args.target_tensor_model_parallel_size):
+        output_state_dict.append({})
+
+    # Embedding layer
+    print("converting embedding layer")
+    pos_embedding = state_dict["transformer.wpe.weight"].to(dtype)
+    word_embedding = state_dict["transformer.wte.weight"].to(dtype)
+    orig_vocab_size = config.vocab_size
+    padded_vocab_size = _vocab_size_with_padding(orig_vocab_size, margs)
+    setattr(margs, "padded_vocab_size", padded_vocab_size)
+    # Cut out extra padding we don't need
+    if orig_vocab_size > padded_vocab_size:
+        full_word_embed = word_embedding[0:padded_vocab_size, :]
+    # Expanding embedding to larger size by replicating final entry
+    elif orig_vocab_size < padded_vocab_size:
+        padding_size = padded_vocab_size - orig_vocab_size
+        full_word_embed = torch.cat((word_embedding, word_embedding[-1].unsqueeze(0).expand(padding_size, -1)))
+    # Same size!
+    else:
+        full_word_embed = word_embedding
+
+    # Split into new tensor model parallel sizes
+    out_word_embed = torch.chunk(full_word_embed, args.target_tensor_model_parallel_size, dim=0)
+    for i in range(args.target_tensor_model_parallel_size):
+        pos_emb_dict = get_element_from_dict_by_path(
+            output_state_dict[i], "model.language_model.embedding.position_embeddings"
+        )
+        pos_emb_dict["weight"] = pos_embedding
+
+        word_emb_dict = get_element_from_dict_by_path(
+            output_state_dict[i], "model.language_model.embedding.word_embeddings"
+        )
+        word_emb_dict["weight"] = out_word_embed[i].clone()
+
+    # Transformer layers
+    print("converting transformer layers")
+    if config.num_attention_heads % args.target_tensor_model_parallel_size != 0:
+        raise ValueError(
+            f"Number of attention heads ({config.num_attention_heads}) must be divisible by number of tensor parallelism"
+            f" ({args.target_tensor_model_parallel_size})"
+        )
+
+    if config.num_hidden_layers % args.target_pipeline_model_parallel_size != 0:
+        raise ValueError(
+            f"Number of layers ({config.num_hidden_layers}) must be divisible by number of pipeline parallelism"
+            f" ({args.target_pipeline_model_parallel_size})"
+        )
+
+    num_layers = config.num_hidden_layers // args.target_pipeline_model_parallel_size
+
+    layer_re = re.compile(r"transformer.h\.(\d+)\.([a-z0-9_.]+)\.([a-z]+)")
+    # The number of heads.
+    heads = config.n_head
+    # The hidden_size per head.
+    hidden_size_per_head = config.n_embd // config.n_head
+    for pp_rank in range(args.target_pipeline_model_parallel_size):
+        layer_offset = pp_rank * num_layers
+        if pp_rank > 0:
+            output_state_dict = []
+            for i in range(args.target_tensor_model_parallel_size):
+                output_state_dict.append({})
+
+        for layer in range(num_layers):
+            pp_layer_id = layer + layer_offset
+            layers_to_copy = [
+                layer_name for layer_name in state_dict if layer_name.startswith(f"transformer.h.{pp_layer_id}.")
+            ]
+
+            for layer_name in layers_to_copy:
+                m = layer_re.match(layer_name)
+                # Stop if that's not a layer
+                if m is None:
+                    break
+
+                # The index of the layer.
+                _ = int(m.group(1))
+                # The name of the operation.
+                op_name = m.group(2)
+                # Is it a weight or a bias?
+                weight_or_bias = m.group(3)
+
+                params = state_dict[layer_name].to(dtype)
+                # handle layernorm
+                if op_name.startswith("ln"):
+                    out_name = "input_layernorm" if op_name.endswith("1") else "post_attention_layernorm"
+                    layer_name = f"layers.{layer}.{out_name}.{weight_or_bias}"
+
+                # handle attention K, V, Q weights
+                elif op_name.startswith("attn.c_attn") and weight_or_bias == "weight":
+                    # transformers stores D X (3*D) but Megatron-LM expects (3*D) X D.
+                    params = params.transpose(0, 1).contiguous()
+
+                    params = transformers_to_megatron_fix_query_key_value_ordering(
+                        params,
+                        3.0,
+                        3,
+                        heads,
+                        hidden_size_per_head,
+                    )
+                    layer_name = f"layers.{layer}.self_attention.query_key_value.{weight_or_bias}"
+
+                # handle attention K, V, Q bias
+                elif op_name.startswith("attn.c_attn") and weight_or_bias == "bias":
+                    params = transformers_to_megatron_fix_query_key_value_ordering(
+                        params,
+                        3.0,
+                        3,
+                        heads,
+                        hidden_size_per_head,
+                    )
+                    layer_name = f"layers.{layer}.self_attention.query_key_value.{weight_or_bias}"
+
+                # handle attention and mlp weights
+                elif weight_or_bias == "weight":
+                    out_name = transformers_to_megatron.get(op_name)
+                    if out_name is None:
+                        continue
+                    params = params.transpose(0, 1)
+                    layer_name = f"layers.{layer}.{out_name}.{weight_or_bias}"
+
+                # handle attention and mlp bias
+                elif weight_or_bias == "bias":
+                    out_name = transformers_to_megatron.get(op_name)
+                    if out_name is None:
+                        continue
+                    layer_name = f"layers.{layer}.{out_name}.{weight_or_bias}"
+
+                # skip
+                else:
+                    continue
+
+                if op_name + "." + weight_or_bias in tensor_parallel_params:
+                    dim = 1 if op_name in ["attn.c_proj", "mlp.c_proj"] else 0
+                    params = torch.chunk(params, args.target_tensor_model_parallel_size, dim=dim)
+
+                for i in range(args.target_tensor_model_parallel_size):
+                    params_dict = get_element_from_dict_by_path(output_state_dict[i], "model.language_model.encoder")
+                    params_dict[layer_name] = (
+                        params[i].clone() if (op_name + "." + weight_or_bias in tensor_parallel_params) else params
+                    )
+
+        if pp_rank == args.target_pipeline_model_parallel_size - 1:
+            # handle final layernorm
+            for weight_or_bias in ["weight", "bias"]:
+                params = state_dict[f"transformer.ln_f.{weight_or_bias}"].to(dtype)
+                layer_name = f"final_layernorm.{weight_or_bias}"
+                for i in range(args.target_tensor_model_parallel_size):
+                    params_dict = get_element_from_dict_by_path(output_state_dict[i], "model.language_model.encoder")
+                    params_dict[layer_name] = params
+
+            # add the LM head
+            for i in range(args.target_tensor_model_parallel_size):
+                params_dict = get_element_from_dict_by_path(output_state_dict[i], "model.word_embeddings_for_head")
+                params_dict["weight"] = out_word_embed[i].clone()
+
+        # saving the state dict as per the tp_rank and pp_rank
+        for tp_rank in range(args.target_tensor_model_parallel_size):
+            output_state_dict[tp_rank]["checkpoint_version"] = 3.0
+            output_state_dict[tp_rank]["args"] = margs
+            checkpoint_dir = (
+                f"mp_rank_{tp_rank:02d}"
+                if args.target_pipeline_model_parallel_size == 1
+                else f"mp_rank_{tp_rank:02d}_{pp_rank:03d}"
+            )
+            if args.use_distributed_optimizer:
+                checkpoint_name = "model_rng.pt"
+            else:
+                checkpoint_name = "model_optim_rng.pt"
+                output_state_dict[tp_rank]["optimizer"] = dummy_optim_state_dict["optimizer"]
+            checkpoint_dir = os.path.join(release_dir, checkpoint_dir)
+            os.makedirs(checkpoint_dir, exist_ok=True)
+            checkpoint_path = os.path.join(checkpoint_dir, checkpoint_name)
+            if args.print_checkpoint_structure:
+                print(
+                    f"Checkpoint structure of model state dict shard belonging to TP rank {tp_rank} and PP rank"
+                    f" {pp_rank}:"
+                )
+                recursive_print(None, output_state_dict[tp_rank])
+            torch.save(output_state_dict[tp_rank], checkpoint_path)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser = add_checkpointing_args(parser)
+    parser = add_megatron_checkpoint_args(parser)
+    parser = add_transformers_checkpoint_args(parser)
+    args = parser.parse_args()
+    if args.convert_checkpoint_from_megatron_to_transformers:
+        convert_checkpoint_from_megatron_to_transformers(args)
+    else:
+        convert_checkpoint_from_transformers_to_megatron(args)
+
+
+if __name__ == "__main__":
+    main()

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/mllama/configuration_mllama.py

@@ -0,0 +1,200 @@
+# Copyright 2024 HuggingFace Inc. team. All rights reserved.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Mllama model configuration"""
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring, logging
+
+
+logger = logging.get_logger(__name__)
+
+
+@auto_docstring(checkpoint="meta-llama/Llama-3.2-11B-Vision")
+@strict
+class MllamaVisionConfig(PreTrainedConfig):
+    r"""
+    num_global_layers (`int`, *optional*, defaults to 8):
+        Number of global layers in the Transformer encoder. Vision model has a second transformer encoder, called global.
+    vision_output_dim (`int`, *optional*, defaults to 7680):
+        Dimensionality of the vision model output. Includes output of transformer
+        encoder with intermediate layers and global transformer encoder.
+    max_num_tiles (`int`, *optional*, defaults to 4):
+        Maximum number of tiles for image splitting.
+    intermediate_layers_indices (`list[int]`, *optional*, defaults to [3, 7, 15, 23, 30]):
+        Indices of intermediate layers of transformer encoder from which to extract and output features.
+        These output features are concatenated with final hidden state of transformer encoder.
+    supported_aspect_ratios (`list[list[int]]`, *optional*):
+        List of supported aspect ratios for image splitting. If not specified, the default supported aspect ratios
+        are [[1, 1], [1, 2], [1, 3], [1, 4], [2, 1], [2, 2], [3, 1], [4, 1]] for `max_num_tiles=4`.
+
+    Example:
+
+    ```python
+    >>> from transformers import MllamaVisionConfig, MllamaVisionModel
+
+    >>> # Initializing a Llama config
+    >>> config = MllamaVisionConfig()
+
+    >>> # Initializing a vision model from the mllama-11b style configuration
+    >>> model = MllamaVisionModel(config)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "mllama_vision_model"
+    base_config_key = "vision_config"
+    attribute_map = {"num_attention_heads": "attention_heads"}
+
+    hidden_size: int = 1280
+    hidden_act: str = "gelu"
+    num_hidden_layers: int = 32
+    num_global_layers: int = 8
+    attention_heads: int = 16
+    num_channels: int = 3
+    intermediate_size: int = 5120
+    vision_output_dim: int = 7680
+    image_size: int | list[int] | tuple[int, int] = 448
+    patch_size: int | list[int] | tuple[int, int] = 14
+    norm_eps: float = 1e-5
+    max_num_tiles: int = 4
+    intermediate_layers_indices: list[int] | None = None
+    supported_aspect_ratios: list[list[int]] | None = None
+    initializer_range: float = 0.02
+
+    def __post_init__(self, **kwargs):
+        if self.supported_aspect_ratios is None:
+            self.supported_aspect_ratios = [[1, 1], [1, 2], [1, 3], [1, 4], [2, 1], [2, 2], [3, 1], [4, 1]]
+
+        if self.intermediate_layers_indices is None:
+            self.intermediate_layers_indices = [3, 7, 15, 23, 30]
+        super().__post_init__(**kwargs)
+
+    def validate_architecture(self):
+        """Part of `@strict`-powered validation. Validates the architecture of the config."""
+        if (
+            self.supported_aspect_ratios == [[1, 1], [1, 2], [1, 3], [1, 4], [2, 1], [2, 2], [3, 1], [4, 1]]
+            and self.max_num_tiles != 4
+        ):
+            raise ValueError("max_num_tiles must be 4 for default supported aspect ratios")
+
+    @property
+    def max_aspect_ratio_id(self) -> int:
+        return len(self.supported_aspect_ratios)
+
+
+@auto_docstring(checkpoint="meta-llama/Llama-3.2-11B-Vision")
+@strict
+class MllamaTextConfig(PreTrainedConfig):
+    r"""
+    cross_attention_layers (`list[int]`, *optional*):
+        Indices of the cross attention layers. If not specified, will default to [3, 8, 13, 18, 23, 28, 33, 38].
+
+    Example:
+
+    ```python
+    >>> from transformers import MllamaTextModel, MllamaTextConfig
+
+    >>> # Initializing a Mllama text config
+    >>> config = MllamaTextConfig()
+
+    >>> # Initializing a model from the Mllama text configuration
+    >>> model = MllamaTextModel(config)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "mllama_text_model"
+    base_config_key = "text_config"
+    default_theta = 500000.0
+
+    vocab_size: int = 128256
+    hidden_size: int = 4096
+    hidden_act: str = "silu"
+    num_hidden_layers: int = 40
+    num_attention_heads: int = 32
+    num_key_value_heads: int = 8
+    intermediate_size: int = 14_336
+    rope_parameters: dict | None = None
+    rms_norm_eps: float = 1e-5
+    max_position_embeddings: int = 131_072
+    initializer_range: float = 0.02
+    use_cache: bool = True
+    tie_word_embeddings: bool = False
+    cross_attention_layers: list[int] | None = None
+    dropout: float | int = 0.0
+    bos_token_id: int = 128000
+    eos_token_id: int | list[int] | None = 128001
+    pad_token_id: int | None = 128004
+
+    def __post_init__(self, **kwargs):
+        if self.cross_attention_layers is None:
+            self.cross_attention_layers = [3, 8, 13, 18, 23, 28, 33, 38]
+        super().__post_init__(**kwargs)
+
+
+@auto_docstring(checkpoint="meta-llama/Llama-3.2-11B-Vision")
+@strict
+class MllamaConfig(PreTrainedConfig):
+    r"""
+    Example:
+
+    ```python
+    >>> from transformers import MllamaForConditionalGeneration, MllamaConfig, MllamaVisionConfig, MllamaTextConfig
+
+    >>> # Initializing a CLIP-vision config
+    >>> vision_config = MllamaVisionConfig()
+
+    >>> # Initializing a Llama config
+    >>> text_config = MllamaTextConfig()
+
+    >>> # Initializing a mllama-11b style configuration
+    >>> configuration = MllamaConfig(vision_config, text_config)
+
+    >>> # Initializing a model from the mllama-11b style configuration
+    >>> model = MllamaForConditionalGeneration(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "mllama"
+    attribute_map = {
+        "image_token_id": "image_token_index",
+    }
+    sub_configs = {"text_config": MllamaTextConfig, "vision_config": MllamaVisionConfig}
+
+    vision_config: dict | PreTrainedConfig | None = None
+    text_config: dict | PreTrainedConfig | None = None
+    image_token_index: int = 128256
+
+    def __post_init__(self, **kwargs):
+        if self.vision_config is None:
+            self.vision_config = MllamaVisionConfig()
+            logger.info("vision_config is None, using default mllama vision config")
+        elif isinstance(self.vision_config, dict):
+            self.vision_config = MllamaVisionConfig(**self.vision_config)
+
+        if self.text_config is None:
+            self.text_config = MllamaTextConfig()
+            logger.info("text_config is None, using default mllama text config")
+        elif isinstance(self.text_config, dict):
+            self.text_config = MllamaTextConfig(**self.text_config)
+
+        super().__post_init__(**kwargs)
+
+
+__all__ = ["MllamaConfig", "MllamaTextConfig", "MllamaVisionConfig"]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/mllama/modeling_mllama.py

@@ -0,0 +1,1622 @@
+# Copyright 2024 the HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PyTorch Mllama model."""
+
+import math
+from collections.abc import Callable
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache
+from ...generation import GenerationMixin
+from ...masking_utils import create_causal_mask
+from ...modeling_flash_attention_utils import FlashAttentionKwargs
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPast, CausalLMOutputWithPast
+from ...modeling_rope_utils import (
+    ROPE_INIT_FUNCTIONS,
+    dynamic_rope_update,
+)
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging
+from ...utils.generic import (
+    maybe_autocast,
+    merge_with_config_defaults,
+)
+from ...utils.output_capturing import OutputRecorder, capture_outputs
+from .configuration_mllama import MllamaConfig, MllamaTextConfig, MllamaVisionConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+def _prepare_cross_attention_mask(
+    cross_attention_mask: torch.Tensor,
+    num_vision_tokens: int,
+    dtype: str,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    # reshape so it can be used by attn module
+    batch_size, text_total_length, *_ = cross_attention_mask.shape
+    cross_attention_mask = cross_attention_mask.repeat_interleave(num_vision_tokens, dim=3)
+    cross_attention_mask = cross_attention_mask.view(batch_size, text_total_length, -1)
+    cross_attention_mask = cross_attention_mask.unsqueeze(1)
+
+    # invert the mask
+    inverted_cross_attn_mask = (1.0 - cross_attention_mask).to(dtype)
+    cross_attention_mask = inverted_cross_attn_mask.masked_fill(
+        inverted_cross_attn_mask.to(torch.bool), torch.finfo(dtype).min
+    )
+
+    # apply full-row bias, which return 4D tensor of shape [B, H, S1, 1] where value is 0 if the a full row in cross attn mask's
+    # last dimension contains negative infinity values, otherwise it's 1
+    negative_inf_value = torch.finfo(dtype).min
+    full_text_row_masked_out_mask = (
+        (cross_attention_mask != negative_inf_value).any(dim=-1).type_as(cross_attention_mask)[..., None]
+    )
+    cross_attention_mask *= full_text_row_masked_out_mask
+
+    return cross_attention_mask, full_text_row_masked_out_mask
+
+
+def _prepare_aspect_ratio_attention_mask(
+    aspect_ratio_mask: torch.Tensor,
+    num_patches: int,
+    target_length: int,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    # Expand aspect ratio mask to target_length
+    batch_size, max_num_tiles = aspect_ratio_mask.shape
+    attention_mask = aspect_ratio_mask.view(batch_size, max_num_tiles, 1, 1).to(dtype)
+    attention_mask = attention_mask.repeat(1, 1, target_length, 1)
+
+    # Mask padding patches
+    pad_patches = target_length - num_patches
+    attention_mask[:, :, -pad_patches:] = 0
+
+    # Invert the mask (0 -> 1, 1 -> 0)
+    attention_mask = 1 - attention_mask
+
+    # Reshape to 2D and create 4D attention mask
+    # (batch_size, 1, max_num_tiles * target_length, max_num_tiles * target_length)
+    attention_mask = attention_mask.reshape(batch_size, max_num_tiles * target_length, 1)
+    attention_mask = attention_mask @ attention_mask.transpose(-1, -2) * torch.finfo(dtype).min
+    attention_mask = attention_mask.unsqueeze(1)
+
+    return attention_mask
+
+
+class MllamaPrecomputedAspectRatioEmbedding(nn.Module):
+    def __init__(self, config: MllamaVisionConfig, is_gated: bool = True):
+        super().__init__()
+        self.max_num_tiles = config.max_num_tiles
+        self.hidden_size = config.hidden_size
+        self.max_aspect_ratio_id = config.max_aspect_ratio_id
+        self.is_gated = is_gated
+
+        self.embedding = nn.Embedding(self.max_aspect_ratio_id + 1, self.max_num_tiles * self.hidden_size)
+        if is_gated:
+            self.gate = nn.Parameter(torch.zeros(1))
+
+    def forward(self, hidden_state: torch.Tensor, aspect_ratio_ids: torch.Tensor) -> torch.Tensor:
+        embeddings = self.embedding(aspect_ratio_ids)
+        embeddings = embeddings.reshape(-1, self.max_num_tiles, 1, self.hidden_size)
+
+        if self.is_gated:
+            embeddings = embeddings * self.gate.tanh()
+
+        hidden_state = hidden_state + embeddings
+        return hidden_state
+
+
+class MllamaPrecomputedPositionEmbedding(nn.Module):
+    def __init__(self, config: MllamaVisionConfig):
+        super().__init__()
+        self.max_num_tiles = config.max_num_tiles
+        self.max_aspect_ratio_id = config.max_aspect_ratio_id
+        self.num_patches = (config.image_size // config.patch_size) ** 2 + 1
+        self.hidden_size = config.hidden_size
+        self.scale = config.hidden_size**-0.5
+
+        self.gate = nn.Parameter(torch.zeros(1))
+
+        # position embedding
+        position_embedding = torch.randn(self.num_patches, self.hidden_size)
+        self.embedding = nn.Parameter(self.scale * position_embedding)
+
+        # tile position embedding
+        self.tile_embedding = nn.Embedding(
+            self.max_aspect_ratio_id + 1, self.max_num_tiles * self.num_patches * self.hidden_size
+        )
+
+    def forward(self, hidden_state: torch.Tensor, aspect_ratio_ids: torch.Tensor) -> torch.Tensor:
+        # position embeddings
+        gated_position_embedding = (1 - self.gate.tanh()) * self.embedding
+        hidden_state = hidden_state + gated_position_embedding.view(1, 1, self.num_patches, self.hidden_size)
+
+        # precomputed tile position embeddings
+        tile_position_embedding = self.tile_embedding(aspect_ratio_ids)
+        batch_size = hidden_state.shape[0]
+        tile_position_embedding = tile_position_embedding.reshape(
+            batch_size, self.max_num_tiles, self.num_patches, self.hidden_size
+        )
+        gated_tile_position_embedding = self.gate.tanh() * tile_position_embedding
+        hidden_state = hidden_state + gated_tile_position_embedding
+
+        return hidden_state
+
+
+# Copied from transformers.models.clip.modeling_clip.CLIPMLP with CLIP->MllamaVision
+class MllamaVisionMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.activation_fn = ACT2FN[config.hidden_act]
+        self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
+        self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.fc1(hidden_states)
+        hidden_states = self.activation_fn(hidden_states)
+        hidden_states = self.fc2(hidden_states)
+        return hidden_states
+
+
+# Copied from transformers.models.llama.modeling_llama.repeat_kv
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+# Copied from transformers.models.llama.modeling_llama.eager_attention_forward
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: torch.Tensor | None,
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs: Unpack[TransformersKwargs],
+):
+    key_states = repeat_kv(key, module.num_key_value_groups)
+    value_states = repeat_kv(value, module.num_key_value_groups)
+
+    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    if attention_mask is not None:
+        attn_weights = attn_weights + attention_mask
+
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+    attn_output = torch.matmul(attn_weights, value_states)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+
+    return attn_output, attn_weights
+
+
+class MllamaVisionAttention(nn.Module):
+    def __init__(self, config: MllamaVisionConfig):
+        super().__init__()
+
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.num_heads = config.attention_heads
+        self.head_dim = config.hidden_size // config.attention_heads
+        self.scaling = self.head_dim**-0.5
+        self.num_key_value_groups = 1
+
+        self.q_proj = nn.Linear(self.embed_dim, self.num_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(self.embed_dim, self.num_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(self.embed_dim, self.num_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.embed_dim, bias=False)
+
+    def forward(
+        self,
+        hidden_state: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        **kwargs,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        query = self.q_proj(hidden_state)
+        key = self.k_proj(hidden_state)
+        value = self.v_proj(hidden_state)
+
+        batch_size, q_seq_len, _ = query.shape
+        _, kv_seq_len, _ = key.shape
+
+        query = query.view(batch_size, q_seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key = key.view(batch_size, kv_seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+        value = value.view(batch_size, kv_seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+
+        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, eager_attention_forward
+        )
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query,
+            key,
+            value,
+            attention_mask,
+            dropout=0.0,
+            scaling=self.scaling,
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(batch_size, q_seq_len, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+
+        return attn_output, attn_weights
+
+
+class MllamaVisionEncoderLayer(nn.Module):
+    def __init__(self, config: MllamaVisionConfig, is_gated: bool = False):
+        super().__init__()
+
+        self.hidden_size = config.hidden_size
+        self.num_attention_heads = config.attention_heads
+        self.is_gated = is_gated
+        self.intermediate_size = config.intermediate_size
+
+        self.self_attn = MllamaVisionAttention(config)
+        self.mlp = MllamaVisionMLP(config)
+
+        self.input_layernorm = nn.LayerNorm(self.hidden_size, eps=config.norm_eps)
+        self.post_attention_layernorm = nn.LayerNorm(self.hidden_size, eps=config.norm_eps)
+
+        if is_gated:
+            self.gate_attn = nn.Parameter(torch.ones(1) * math.pi / 4)
+            self.gate_ffn = nn.Parameter(torch.ones(1) * math.pi / 4)
+
+    def forward(
+        self,
+        hidden_state: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ):
+        # Self Attention
+        residual = hidden_state
+        hidden_state = self.input_layernorm(hidden_state)
+        hidden_state, attn_weights = self.self_attn(hidden_state, attention_mask=attention_mask)
+        if self.is_gated:
+            hidden_state = self.gate_attn.tanh() * hidden_state
+        hidden_state = residual + hidden_state
+
+        # Feed forward
+        residual = hidden_state
+        hidden_state = self.post_attention_layernorm(hidden_state)
+        hidden_state = self.mlp(hidden_state)
+        if self.is_gated:
+            hidden_state = self.gate_ffn.tanh() * hidden_state
+        hidden_state = residual + hidden_state
+
+        return hidden_state
+
+
+class MllamaVisionEncoder(nn.Module):
+    """
+    Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a
+    [`MllamaEncoderLayer`].
+
+    Args:
+        config: MllamaConfig
+    """
+
+    def __init__(self, config: MllamaVisionConfig, num_layers=32, is_gated=False):
+        super().__init__()
+        self.config = config
+        self.layers = nn.ModuleList([MllamaVisionEncoderLayer(config, is_gated) for _ in range(num_layers)])
+        self.gradient_checkpointing = False
+        self.config = config
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> BaseModelOutput:
+        r"""
+        Args:
+            inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+                Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation.
+                This is useful if you want more control over how to convert `input_ids` indices into associated vectors
+                than the model's internal embedding lookup matrix.
+            attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+
+                - 1 for tokens that are **not masked**,
+                - 0 for tokens that are **masked**.
+
+                [What are attention masks?](../glossary#attention-mask)
+
+        """
+        encoder_states = ()
+        for encoder_layer in self.layers:
+            hidden_states = encoder_layer(
+                hidden_state=hidden_states,
+                attention_mask=attention_mask,
+            )
+            encoder_states = encoder_states + (hidden_states,)
+
+        return BaseModelOutput(last_hidden_state=hidden_states, hidden_states=encoder_states)
+
+
+# Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->MllamaText
+class MllamaTextRMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps: float = 1e-6) -> None:
+        """
+        MllamaTextRMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+    def extra_repr(self):
+        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+class MllamaTextCrossAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(
+        self,
+        config: MllamaTextConfig | None = None,
+        layer_idx: int | None = None,
+    ):
+        super().__init__()
+        self.config = config
+        self.num_heads = self.config.num_attention_heads
+        self.num_key_value_heads = self.config.num_key_value_heads
+        self.dropout = config.dropout
+        self.hidden_size = config.hidden_size
+        self.head_dim = config.hidden_size // self.num_heads
+        self.layer_idx = layer_idx
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.scaling = self.head_dim**-0.5
+
+        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
+
+        self.q_norm = MllamaTextRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.k_norm = MllamaTextRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cross_attention_states: torch.Tensor | None = None,
+        past_key_values: Cache | None = None,
+        attention_mask: torch.Tensor | None = None,
+        use_cache: bool | None = None,
+        **kwargs,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]:
+        """Input shape: Batch x Time x Channel"""
+        bsz, q_len, _ = hidden_states.size()
+        query_states = self.q_proj(hidden_states)
+        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        query_states = self.q_norm(query_states)
+
+        if cross_attention_states is not None:
+            key_states = self.k_proj(cross_attention_states)
+            value_states = self.v_proj(cross_attention_states)
+            key_states = key_states.view(bsz, -1, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+            value_states = value_states.view(bsz, -1, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+            key_states = self.k_norm(key_states)
+            if past_key_values is not None:
+                # if we have a new image + new tokens, we only computed key_states on that new image
+                # we still update the cross key states, past_image, new_image. And use it!
+                key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+        elif past_key_values is not None and past_key_values.get_seq_length() > 0:
+            key_states, value_states = (
+                past_key_values.layers[self.layer_idx].keys,
+                past_key_values.layers[self.layer_idx].values,
+            )
+        else:
+            raise ValueError(
+                "Cross attention layer can't find neither `cross_attn_states` nor cached values for key/values!"
+            )
+
+        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, eager_attention_forward
+        )
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0 if not self.training else self.dropout,
+            scaling=self.scaling,
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(bsz, q_len, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+
+        return attn_output, attn_weights
+
+
+# Copied from transformers.models.llama.modeling_llama.rotate_half
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+# Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb
+def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
+    """Applies Rotary Position Embedding to the query and key tensors.
+
+    Args:
+        q (`torch.Tensor`): The query tensor.
+        k (`torch.Tensor`): The key tensor.
+        cos (`torch.Tensor`): The cosine part of the rotary embedding.
+        sin (`torch.Tensor`): The sine part of the rotary embedding.
+        unsqueeze_dim (`int`, *optional*, defaults to 1):
+            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+    Returns:
+        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+    """
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+class MllamaTextSelfAttention(nn.Module):
+    def __init__(self, config: MllamaTextConfig, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.num_heads = config.num_attention_heads
+        self.dropout = config.dropout
+        self.hidden_size = config.hidden_size
+        self.num_key_value_heads = config.num_key_value_heads
+        self.head_dim = config.hidden_size // self.num_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.scaling = self.head_dim**-0.5
+
+        self.layer_idx = layer_idx
+        self.is_causal = True
+
+        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        position_embeddings: torch.Tensor,
+        past_key_values=None,
+        **kwargs,
+    ):
+        bsz, q_len, _ = hidden_states.size()
+
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        if past_key_values is not None:
+            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+
+        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, eager_attention_forward
+        )
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0 if not self.training else self.dropout,
+            scaling=self.scaling,
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(bsz, q_len, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+
+        return attn_output, attn_weights
+
+
+# Copied from transformers.models.gemma2.modeling_gemma2.Gemma2MLP with Gemma2->MllamaText
+class MllamaTextMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        # Ignore copy
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, x):
+        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+        return down_proj
+
+
+# Modified from transformers.models.llama.modeling_llama.LlamaDecoderLayer
+class MllamaSelfAttentionDecoderLayer(GradientCheckpointingLayer):
+    def __init__(self, config: MllamaTextConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+
+        self.self_attn = MllamaTextSelfAttention(config=config, layer_idx=layer_idx)
+
+        self.mlp = MllamaTextMLP(config)
+        self.input_layernorm = MllamaTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = MllamaTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        self.layer_idx = layer_idx
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cross_attention_states: torch.Tensor | None = None,
+        cross_attention_mask: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        full_text_row_masked_out_mask: tuple[torch.Tensor, torch.Tensor] | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        use_cache: bool | None = False,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]:
+        """
+        Args:
+            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+            attention_mask (`torch.FloatTensor`, *optional*):
+                attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1,
+                query_sequence_length, key_sequence_length)` if default attention is used.
+
+            use_cache (`bool`, *optional*):
+                If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
+                (see `past_key_values`).
+            past_key_values (`Cache`, *optional*): cached past key and value projection states
+            position_embeddings (`tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*):
+                Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`,
+                with `head_dim` being the embedding dimension of each attention head.
+            kwargs (`dict`, *optional*):
+                Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code
+                into the model
+        """
+        residual = hidden_states
+
+        hidden_states = self.input_layernorm(hidden_states)
+
+        # Self Attention
+        hidden_states, self_attn_weights = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        return hidden_states
+
+
+class MllamaCrossAttentionDecoderLayer(GradientCheckpointingLayer):
+    """Cross-attention transformer block with tanh-gated attention and feedforward."""
+
+    def __init__(self, config: MllamaTextConfig, layer_idx: int) -> None:
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.cross_attn = MllamaTextCrossAttention(config, layer_idx=layer_idx)
+
+        self.input_layernorm = MllamaTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.cross_attn_attn_gate = torch.nn.Parameter(torch.zeros(1))
+
+        self.mlp = MllamaTextMLP(config)
+        self.post_attention_layernorm = MllamaTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.cross_attn_mlp_gate = torch.nn.Parameter(torch.zeros(1))
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cross_attention_states: torch.Tensor,
+        cross_attention_mask: torch.Tensor,
+        attention_mask: torch.Tensor,
+        full_text_row_masked_out_mask: tuple[torch.Tensor, torch.Tensor],
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        use_cache: bool | None = False,
+        position_embeddings: torch.Tensor | None = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> tuple[torch.Tensor]:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+
+        hidden_states, attn_weights = self.cross_attn(
+            hidden_states=hidden_states,
+            attention_mask=cross_attention_mask,
+            cross_attention_states=cross_attention_states,
+            past_key_values=past_key_values,
+            **kwargs,
+        )
+        hidden_states = residual + self.cross_attn_attn_gate.tanh() * hidden_states
+
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        if full_text_row_masked_out_mask is not None:
+            hidden_states = full_text_row_masked_out_mask[:, 0] * hidden_states  # type: ignore
+        hidden_states = residual + self.cross_attn_mlp_gate.tanh() * hidden_states
+
+        return hidden_states
+
+
+# Copied from transformers.models.llama.modeling_llama.LlamaRotaryEmbedding with LlamaConfig->MllamaTextConfig,Llama->Mllama
+class MllamaRotaryEmbedding(nn.Module):
+    inv_freq: torch.Tensor  # fix linting for `register_buffer`
+
+    def __init__(self, config: MllamaTextConfig, device=None):
+        super().__init__()
+        self.max_seq_len_cached = config.max_position_embeddings
+        self.original_max_seq_len = config.max_position_embeddings
+
+        self.config = config
+
+        self.rope_type = self.config.rope_parameters["rope_type"]
+        rope_init_fn: Callable = self.compute_default_rope_parameters
+        if self.rope_type != "default":
+            rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+        inv_freq, self.attention_scaling = rope_init_fn(self.config, device)
+
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False)
+
+    @staticmethod
+    def compute_default_rope_parameters(
+        config: MllamaTextConfig | None = None,
+        device: Optional["torch.device"] = None,
+        seq_len: int | None = None,
+    ) -> tuple["torch.Tensor", float]:
+        """
+        Computes the inverse frequencies according to the original RoPE implementation
+        Args:
+            config ([`~transformers.PreTrainedConfig`]):
+                The model configuration.
+            device (`torch.device`):
+                The device to use for initialization of the inverse frequencies.
+            seq_len (`int`, *optional*):
+                The current sequence length. Unused for this type of RoPE.
+        Returns:
+            Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
+            post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
+        """
+        base = config.rope_parameters["rope_theta"]
+        dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+
+        attention_factor = 1.0  # Unused in this type of RoPE
+
+        # Compute the inverse frequencies
+        inv_freq = 1.0 / (
+            base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
+        )
+        return inv_freq, attention_factor
+
+    # Ignore copy
+    @torch.no_grad()
+    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
+    def forward(self, x, position_ids):
+        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
+        position_ids_expanded = position_ids[:, None, :].float()
+
+        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+        with maybe_autocast(device_type=device_type, enabled=False):  # Force float32
+            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos() * self.attention_scaling
+            sin = emb.sin() * self.attention_scaling
+
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+@auto_docstring
+class MllamaPreTrainedModel(PreTrainedModel):
+    config: MllamaConfig
+    base_model_prefix = "model"
+    input_modalities = ("image", "text")
+    supports_gradient_checkpointing = True
+    _no_split_modules = [
+        "MllamaVisionEncoderLayer",
+        "MllamaCrossAttentionDecoderLayer",
+        "MllamaSelfAttentionDecoderLayer",
+    ]
+    _can_compile_fullgraph = False  # static cache cannot have different shapes for each layer
+    _supports_sdpa = True
+    _supports_flash_attn = True
+    _supports_flex_attn = True
+    _supports_attention_backend = True
+    _can_record_outputs = {
+        "hidden_states": [MllamaSelfAttentionDecoderLayer, MllamaCrossAttentionDecoderLayer],
+        "attentions": [
+            OutputRecorder(MllamaTextSelfAttention, index=1, layer_name="self_attn"),
+            OutputRecorder(MllamaTextSelfAttention, index=1, layer_name="cross_attn"),
+            OutputRecorder(MllamaTextCrossAttention, index=1, layer_name="cross_attn"),
+        ],
+    }
+
+    @torch.no_grad()
+    def _init_weights(self, module):
+        std = getattr(self.config, "initializer_range", self.config.get_text_config().initializer_range)
+
+        if isinstance(module, (nn.Linear, nn.Conv2d)):
+            init.normal_(module.weight, mean=0.0, std=std)
+            if module.bias is not None:
+                init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            init.normal_(module.weight, mean=0.0, std=std)
+            # Here we need the check explicitly, as we slice the weight in the `zeros_` call, so it looses the flag
+            if module.padding_idx is not None and not getattr(module.weight, "_is_hf_initialized", False):
+                init.zeros_(module.weight[module.padding_idx])
+        elif isinstance(module, nn.LayerNorm):
+            init.ones_(module.weight)
+            init.zeros_(module.bias)
+        elif isinstance(module, MllamaTextRMSNorm):
+            init.ones_(module.weight)
+        elif isinstance(module, MllamaVisionModel):
+            init.normal_(module.class_embedding, std=std)
+        elif isinstance(module, MllamaPrecomputedPositionEmbedding):
+            init.normal_(module.embedding, std=std)
+            init.zeros_(module.gate)
+        elif isinstance(module, MllamaVisionEncoderLayer) and module.is_gated:
+            init.normal_(module.gate_attn, std=std)
+            init.normal_(module.gate_ffn, std=std)
+        elif isinstance(module, MllamaCrossAttentionDecoderLayer):
+            init.zeros_(module.cross_attn_attn_gate)
+            init.zeros_(module.cross_attn_mlp_gate)
+        elif isinstance(module, MllamaPrecomputedAspectRatioEmbedding):
+            if module.is_gated:
+                init.zeros_(module.gate)
+        elif isinstance(module, MllamaRotaryEmbedding):
+            rope_fn = (
+                ROPE_INIT_FUNCTIONS[module.rope_type]
+                if module.rope_type != "default"
+                else module.compute_default_rope_parameters
+            )
+            buffer_value, _ = rope_fn(module.config)
+            init.copy_(module.inv_freq, buffer_value)
+            init.copy_(module.original_inv_freq, buffer_value)
+
+
+@auto_docstring(
+    custom_intro="""
+    The Mllama Vision Model which consists of two vision encoders.
+    """
+)
+class MllamaVisionModel(MllamaPreTrainedModel):
+    config: MllamaVisionConfig
+    base_model_prefix = "vision_model"
+    input_modalities = ("image",)
+
+    def __init__(self, config: MllamaVisionConfig):
+        super().__init__(config)
+        self.image_size = config.image_size
+        self.patch_size = config.patch_size
+        self.max_num_tiles = config.max_num_tiles
+        self.hidden_size = config.hidden_size
+        self.num_channels = config.num_channels
+        self.intermediate_layers_indices = config.intermediate_layers_indices
+
+        self.num_patches = (self.image_size // self.patch_size) ** 2 + 1
+        self.scale = config.hidden_size**-0.5
+
+        self.patch_embedding = nn.Conv2d(
+            in_channels=config.num_channels,
+            out_channels=self.hidden_size,
+            kernel_size=self.patch_size,
+            stride=self.patch_size,
+            padding="valid",
+            bias=False,
+        )
+
+        self.class_embedding = nn.Parameter(self.scale * torch.randn(self.hidden_size))
+        self.gated_positional_embedding = MllamaPrecomputedPositionEmbedding(config)
+
+        self.pre_tile_positional_embedding = MllamaPrecomputedAspectRatioEmbedding(config, is_gated=True)
+        self.post_tile_positional_embedding = MllamaPrecomputedAspectRatioEmbedding(config, is_gated=True)
+
+        # layer norms
+        self.layernorm_pre = nn.LayerNorm(self.hidden_size)
+        self.layernorm_post = nn.LayerNorm(self.hidden_size)
+
+        # encoders
+        self.transformer = MllamaVisionEncoder(config, config.num_hidden_layers, is_gated=False)
+        self.global_transformer = MllamaVisionEncoder(config, config.num_global_layers, is_gated=True)
+
+        self.post_init()
+
+    def get_input_embeddings(self):
+        """
+        This function is used to fetch the first embedding layer to activate grads on inputs.
+        """
+        return self.patch_embedding
+
+    def apply_class_embedding(self, hidden_state: torch.Tensor) -> torch.Tensor:
+        batch_size, _, hidden_size = hidden_state.shape
+        class_embedding = self.class_embedding.expand(batch_size, 1, hidden_size)
+        hidden_state = torch.cat([class_embedding, hidden_state], dim=1)
+        return hidden_state
+
+    @merge_with_config_defaults
+    @capture_outputs
+    @auto_docstring
+    def forward(
+        self, pixel_values: torch.Tensor, aspect_ratio_ids: torch.Tensor, aspect_ratio_mask: torch.Tensor, **kwargs
+    ) -> BaseModelOutput:
+        r"""
+        aspect_ratio_ids (`torch.Tensor` of shape `(batch_size, max_num_images)`, *optional*):
+            Aspect ratio ids used to select the appropriate precomputed tile embeddings based on the aspect ratio of each input image.
+            These ids correspond to indices in the model's list of supported aspect ratios, offset by 1.
+
+            For example, if the model supports aspect ratios [[1, 1], [1, 2], [2, 1]]:
+            - An image with aspect ratio [1, 1] would have ID 1
+            - An image with aspect ratio [1, 2] would have ID 2
+            - An image with aspect ratio [2, 1] would have ID 3
+
+            The id 0 is reserved for padding (i.e., no image).
+
+            If an image has aspect ratio [1, 2], that means it was split into 2 tiles horizontally, and its `aspect_ratio_id` would be 2.
+        aspect_ratio_mask (`torch.Tensor` of shape `(batch_size, max_num_images, max_num_tiles)`, *optional*):
+            Mask to avoid performing attention on padding tiles. Mask values selected in `[0, 1]`:
+
+            - 1 for tiles that are **not masked**,
+            - 0 for tiles that are **masked**.
+
+        Example:
+
+        ```python
+        >>> from PIL import Image
+        >>> import httpx
+        >>> from io import BytesIO
+        >>> from transformers import AutoProcessor, MllamaVisionModel
+
+        >>> checkpoint = "meta-llama/Llama-3.2-11B-Vision"
+        >>> model = MllamaVisionModel.from_pretrained(checkpoint)
+        >>> processor = AutoProcessor.from_pretrained(checkpoint)
+
+        >>> url = "https://www.ilankelman.org/stopsigns/australia.jpg"
+        >>> with httpx.stream("GET", url) as response:
+        ...     image = Image.open(BytesIO(response.read()))
+        >>> inputs = processor(images=image, return_tensors="pt")
+
+        >>> output = model(**inputs)
+
+        >>> print(output.last_hidden_state.shape)
+        torch.Size([1, 1, 4, 1025, 7680])
+        ```
+        """
+        batch_size, num_concurrent_media, num_tiles, num_channels, height, width = pixel_values.shape
+
+        pixel_values = pixel_values.reshape(batch_size * num_concurrent_media * num_tiles, num_channels, height, width)
+        aspect_ratio_ids = aspect_ratio_ids.reshape(batch_size * num_concurrent_media, -1)
+
+        # Patch embedding
+        target_dtype = self.patch_embedding.weight.dtype
+        target_device = self.patch_embedding.weight.device
+        patch_embeds = self.patch_embedding(pixel_values.to(target_device, target_dtype))
+        hidden_state = patch_embeds.flatten(2).transpose(1, 2)
+
+        # Tile embeddings
+        _, num_patches, dim = hidden_state.shape
+        hidden_state = hidden_state.reshape(batch_size * num_concurrent_media, num_tiles, -1, dim)
+        hidden_state = self.pre_tile_positional_embedding(hidden_state, aspect_ratio_ids)
+
+        # Add cls token
+        hidden_state = hidden_state.reshape(batch_size * num_concurrent_media * num_tiles, num_patches, dim)
+        hidden_state = self.apply_class_embedding(hidden_state)
+        num_patches += 1
+
+        # Position embeddings
+        hidden_state = hidden_state.reshape(batch_size * num_concurrent_media, num_tiles, num_patches, dim)
+        hidden_state = self.gated_positional_embedding(hidden_state, aspect_ratio_ids)
+
+        hidden_state = self.layernorm_pre(hidden_state)
+
+        # Compute the number of tokens to pad
+        num_padding_patches = (8 - (hidden_state.shape[-2] % 8)) % 8
+        # Compute padding tuple for pad function
+        padding = (0, 0, 0, num_padding_patches)  # (pad_left, pad_right, pad_left for dim -2, pad_right for dim -2)
+        # Pad the tensor
+        hidden_state = F.pad(hidden_state, padding, mode="constant", value=0)
+        slice_index = -num_padding_patches if num_padding_patches > 0 else None
+
+        # Prepare attention mask
+        attention_mask = aspect_ratio_mask.reshape(batch_size * num_concurrent_media, -1)
+        attention_mask = _prepare_aspect_ratio_attention_mask(
+            aspect_ratio_mask=attention_mask,
+            num_patches=self.num_patches,
+            target_length=hidden_state.shape[2],
+            dtype=self.dtype,
+        )
+
+        # Apply encoder
+        hidden_state = hidden_state.view(batch_size * num_concurrent_media, -1, dim)
+        output = self.transformer(
+            hidden_state,
+            attention_mask=attention_mask,
+        )
+        hidden_state = output.last_hidden_state
+
+        hidden_state = self.layernorm_post(hidden_state)
+
+        # Apply global encoder
+        hidden_state = hidden_state.reshape(
+            batch_size * num_concurrent_media, num_tiles, num_patches + num_padding_patches, dim
+        )
+        hidden_state = self.post_tile_positional_embedding(hidden_state, aspect_ratio_ids)
+        hidden_state = hidden_state.reshape(
+            batch_size * num_concurrent_media, num_tiles * (num_patches + num_padding_patches), dim
+        )
+        global_output = self.global_transformer(
+            hidden_state,
+            attention_mask=attention_mask,
+        )
+        hidden_state = global_output.last_hidden_state
+
+        # Remove padding form hidden state
+        hidden_state = hidden_state.reshape(
+            batch_size * num_concurrent_media, num_tiles, num_patches + num_padding_patches, dim
+        )
+        hidden_state = hidden_state[:, :, :slice_index]
+        hidden_state = hidden_state.reshape(batch_size, num_concurrent_media, num_tiles, num_patches, dim)
+
+        # Collect intermediate layer outputs from encoder output
+        all_intermediate_hidden_states = [output.hidden_states[i] for i in self.intermediate_layers_indices]
+        intermediate_hidden_states = torch.stack(all_intermediate_hidden_states, dim=-1)
+
+        # Remove padding from intermediate hidden states
+        intermediate_hidden_states = intermediate_hidden_states.reshape(
+            batch_size * num_concurrent_media, num_tiles, num_patches + num_padding_patches, -1
+        )
+        intermediate_hidden_states = intermediate_hidden_states[:, :, :slice_index]
+        intermediate_hidden_states = intermediate_hidden_states.reshape(
+            batch_size, num_concurrent_media, num_tiles, num_patches, -1
+        )
+
+        # Concatenate final hidden state and intermediate hidden states
+        hidden_state = torch.cat([hidden_state, intermediate_hidden_states], dim=-1)
+
+        return BaseModelOutput(last_hidden_state=hidden_state)
+
+
+@auto_docstring(
+    custom_intro="""
+    The Mllama Text Model which consists of transformer with self and cross attention layers.
+    """
+)
+class MllamaTextModel(MllamaPreTrainedModel):
+    config: MllamaTextConfig
+    base_model_prefix = "language_model.model"
+    input_modalities = ("text",)
+
+    def __init__(self, config: MllamaTextConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+        self.embed_tokens = nn.Embedding(config.vocab_size + 8, config.hidden_size, self.padding_idx)
+        self.cross_attention_layers = config.cross_attention_layers
+
+        layers = []
+        for layer_idx in range(config.num_hidden_layers):
+            if layer_idx in self.cross_attention_layers:
+                layers.append(MllamaCrossAttentionDecoderLayer(config, layer_idx))
+            else:
+                layers.append(MllamaSelfAttentionDecoderLayer(config, layer_idx))
+
+        self.layers = nn.ModuleList(layers)
+        self.norm = MllamaTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = MllamaRotaryEmbedding(config=config)
+
+        self.gradient_checkpointing = False
+        self.post_init()
+
+    @merge_with_config_defaults
+    @capture_outputs
+    @can_return_tuple
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        cross_attention_states: torch.FloatTensor | None = None,
+        cross_attention_mask: torch.Tensor | None = None,
+        full_text_row_masked_out_mask: tuple[torch.Tensor, torch.Tensor] | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        use_cache: bool | None = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> BaseModelOutputWithPast:
+        r"""
+        cross_attention_states (`torch.FloatTensor`, *optional*):
+            Output of the vision model, used for cross-attention. This tensor contains the processed image features that
+            the language model will attend to.
+        cross_attention_mask (`torch.Tensor` of shape `(batch_size, seq_length, max_num_images, max_num_tiles)`, *optional*):
+            Cross-attention mask to control the interaction between text tokens and image tiles.
+            This 4D tensor defines which image tiles each text token should attend to.
+
+            For each text token (in seq_length):
+            - 1 indicates the token **should attend** to the corresponding image tile
+            - 0 indicates the token **should not attend** to the corresponding image tile
+        full_text_row_masked_out_mask (`tuple[torch.Tensor, torch.Tensor]`, *optional*):
+            A tuple containing two tensors that mask out rows in the cross-attention mechanism:
+            - The first tensor has shape `(batch_size, 1, seq_length, 1)` and contains values of 0 or 1.
+              A value of 0 indicates that the corresponding text token's entire row in the cross-attention
+              matrix should be masked out (all image tokens ignored).
+            - The second tensor has the same shape and is used internally to apply the masking during
+              the forward pass of cross-attention layers.
+            This mask is derived from the cross_attention_mask and is used to handle cases where a text token
+            should not attend to any image token.
+
+        Example:
+
+        ```python
+        >>> from transformers import AutoProcessor, MllamaTextModel
+
+        >>> checkpoint = "meta-llama/Llama-3.2-11B-Vision"
+        >>> model = MllamaTextModel.from_pretrained(checkpoint)
+        >>> processor = AutoProcessor.from_pretrained(checkpoint)
+
+        >>> text = "<|image|>If I had to write a haiku for this one"
+        >>> inputs = processor(text=text, return_tensors="pt")
+
+        >>> output = model(**inputs)
+
+        >>> print(output.last_hidden_state.shape)
+        torch.Size([1, 13, 4096])
+        ```
+        """
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        hidden_states = inputs_embeds
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache(config=self.config)
+
+        if position_ids is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
+            position_ids = position_ids.unsqueeze(0)
+
+        causal_mask = create_causal_mask(
+            config=self.config,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            position_ids=position_ids,
+        )
+        position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids)
+
+        # decoder layers
+        for idx, decoder_layer in enumerate(self.layers):
+            # For text-only path we should skip cross attention layers.
+            # Let's check if the layer is cross attention layer and if we have cross attention states
+            # or cached cross attention states.
+            is_cross_attention_layer = idx in self.cross_attention_layers
+            is_cross_attention_cache_empty = past_key_values is None or (
+                past_key_values is not None and past_key_values.get_seq_length(idx) == 0
+            )
+
+            if is_cross_attention_layer and cross_attention_states is None and is_cross_attention_cache_empty:
+                continue
+
+            hidden_states = decoder_layer(
+                hidden_states,
+                cross_attention_states=cross_attention_states,
+                cross_attention_mask=cross_attention_mask,
+                attention_mask=causal_mask,
+                full_text_row_masked_out_mask=full_text_row_masked_out_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+                position_embeddings=position_embeddings,
+                **kwargs,
+            )
+
+        hidden_states = self.norm(hidden_states)
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+        )
+
+
+@auto_docstring(
+    custom_intro="""
+    The Mllama Text Model with a language modeling head on top.
+    """
+)
+class MllamaForCausalLM(MllamaPreTrainedModel, GenerationMixin):
+    config: MllamaTextConfig
+    _can_compile_fullgraph = True  # only the LLM without cross attn can do compile
+    base_model_prefix = "language_model"
+
+    def __init__(self, config):
+        super().__init__(config.get_text_config())
+        self.text_config = config.get_text_config()
+        self.vocab_size = self.text_config.vocab_size
+        self.model = MllamaTextModel._from_config(self.text_config)
+        self.lm_head = nn.Linear(self.text_config.hidden_size, self.vocab_size, bias=False)
+
+        self.post_init()
+
+    @can_return_tuple
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        cross_attention_states: torch.LongTensor | None = None,
+        cross_attention_mask: torch.LongTensor | None = None,
+        full_text_row_masked_out_mask: tuple[torch.Tensor, torch.Tensor] | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        logits_to_keep: int | torch.Tensor = 0,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple | CausalLMOutputWithPast:
+        r"""
+        cross_attention_states (`torch.FloatTensor`, *optional*):
+            Output of the vision model, used for cross-attention. This tensor contains the processed image features that
+            the language model will attend to.
+        cross_attention_mask (`torch.Tensor` of shape `(batch_size, seq_length, max_num_images, max_num_tiles)`, *optional*):
+            Cross-attention mask to control the interaction between text tokens and image tiles.
+            This 4D tensor defines which image tiles each text token should attend to.
+
+            For each text token (in seq_length):
+            - 1 indicates the token **should attend** to the corresponding image tile
+            - 0 indicates the token **should not attend** to the corresponding image tile
+        full_text_row_masked_out_mask (`tuple[torch.Tensor, torch.Tensor]`, *optional*):
+            A tuple containing two tensors that mask out rows in the cross-attention mechanism:
+            - The first tensor has shape `(batch_size, 1, seq_length, 1)` and contains values of 0 or 1.
+              A value of 0 indicates that the corresponding text token's entire row in the cross-attention
+              matrix should be masked out (all image tokens ignored).
+            - The second tensor has the same shape and is used internally to apply the masking during
+              the forward pass of cross-attention layers.
+            This mask is derived from the cross_attention_mask and is used to handle cases where a text token
+            should not attend to any image token.
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+        Example:
+
+        ```python
+        >>> from transformers import AutoTokenizer, MllamaForCausalLM
+
+        >>> model = MllamaForCausalLM.from_pretrained("Llama-3.2-11B-Vision")
+        >>> tokenizer = AutoTokenizer.from_pretrained("Llama-3.2-11B-Vision")
+
+        >>> prompt = "If I had to write a haiku, it would be:"
+        >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+        >>> # Generate
+        >>> generate_ids = model.generate(inputs.input_ids, max_length=40, do_sample=True, temperature=0.6)
+        >>> result = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        >>> print(result)
+        If I had to write a haiku, it would be: "Snowflakes gently fall" - simple, yet peaceful.
+        I love the idea of snowflakes gently falling, each one
+        ```
+        """
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs = self.model(
+            input_ids=input_ids,
+            cross_attention_states=cross_attention_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            cross_attention_mask=cross_attention_mask,
+            full_text_row_masked_out_mask=full_text_row_masked_out_mask,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            **kwargs,
+        )
+
+        hidden_states = outputs.last_hidden_state
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :]).float()
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(logits, labels, self.vocab_size, **kwargs)
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+@auto_docstring(
+    custom_intro="""
+    The Mllama model which consists of a vision encoder and a language model without language modeling head.
+    """
+)
+class MllamaModel(MllamaPreTrainedModel):
+    def __init__(self, config: MllamaConfig):
+        super().__init__(config)
+        self.vocab_size = config.text_config.vocab_size
+        self.hidden_size = config.text_config.hidden_size
+        self.max_num_tiles = config.vision_config.max_num_tiles
+        self.vision_output_dim = config.vision_config.vision_output_dim
+
+        self.vision_model = MllamaVisionModel._from_config(config.vision_config)
+        self.language_model = MllamaTextModel._from_config(config.text_config)
+        self.multi_modal_projector = nn.Linear(
+            config.vision_config.vision_output_dim,
+            config.text_config.hidden_size,
+            bias=True,
+        )
+        self.post_init()
+
+    @can_return_tuple
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        pixel_values: torch.FloatTensor | None = None,
+        aspect_ratio_mask: torch.Tensor | None = None,
+        aspect_ratio_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        cross_attention_mask: torch.Tensor | None = None,
+        cross_attention_states: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        use_cache: bool | None = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> BaseModelOutputWithPast:
+        r"""
+        aspect_ratio_mask (`torch.Tensor` of shape `(batch_size, max_num_images, max_num_tiles)`, *optional*):
+            Mask to avoid performing attention on padding tiles. Mask values selected in `[0, 1]`:
+
+            - 1 for tiles that are **not masked**,
+            - 0 for tiles that are **masked**.
+        aspect_ratio_ids (`torch.Tensor` of shape `(batch_size, max_num_images)`, *optional*):
+            Aspect ratio ids used to select the appropriate precomputed tile embeddings based on the aspect ratio of each input image.
+            These ids correspond to indices in the model's list of supported aspect ratios, offset by 1.
+
+            For example, if the model supports aspect ratios [[1, 1], [1, 2], [2, 1]]:
+            - An image with aspect ratio [1, 1] would have ID 1
+            - An image with aspect ratio [1, 2] would have ID 2
+            - An image with aspect ratio [2, 1] would have ID 3
+
+            The id 0 is reserved for padding (i.e., no image).
+
+            If an image has aspect ratio [1, 2], that means it was split into 2 tiles horizontally, and its `aspect_ratio_id` would be 2.
+        cross_attention_mask (`torch.Tensor` of shape `(batch_size, seq_length, max_num_images, max_num_tiles)`, *optional*):
+            Cross-attention mask to control the interaction between text tokens and image tiles.
+            This 4D tensor defines which image tiles each text token should attend to.
+
+            For each text token (in seq_length):
+            - 1 indicates the token **should attend** to the corresponding image tile
+            - 0 indicates the token **should not attend** to the corresponding image tile
+        cross_attention_states (`torch.FloatTensor`, *optional*):
+            Output of the vision model, used for cross-attention. This tensor contains the processed image features that
+            the language model will attend to.
+        """
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if pixel_values is not None and cross_attention_states is not None:
+            raise ValueError("`pixel_values` and `cross_attention_states` cannot be provided simultaneously")
+
+        if pixel_values is not None:
+            if aspect_ratio_ids is None:
+                raise ValueError("`aspect_ratio_ids` must be provided if `pixel_values` is provided")
+            # get vision tokens from vision model
+            vision_outputs = self.vision_model(
+                pixel_values=pixel_values,
+                aspect_ratio_ids=aspect_ratio_ids,
+                aspect_ratio_mask=aspect_ratio_mask,
+            )
+            cross_attention_states = vision_outputs.last_hidden_state
+            cross_attention_states = self.multi_modal_projector(cross_attention_states).reshape(
+                -1, cross_attention_states.shape[-2], self.hidden_size
+            )
+
+        if cross_attention_mask is not None:
+            cross_attention_mask, full_text_row_masked_out_mask = _prepare_cross_attention_mask(
+                cross_attention_mask,
+                num_vision_tokens=self.vision_model.num_patches,
+                dtype=self.dtype,
+            )
+        else:
+            full_text_row_masked_out_mask = None
+
+        if cross_attention_mask is not None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            seq_len = input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1]
+            device = input_ids.device if input_ids is not None else inputs_embeds.device
+            current_pos = torch.arange(seq_len, device=device) + past_seen_tokens
+
+            cross_attention_mask = cross_attention_mask[:, :, current_pos]
+            full_text_row_masked_out_mask = full_text_row_masked_out_mask[:, :, current_pos]
+
+        outputs = self.language_model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            cross_attention_states=cross_attention_states,
+            cross_attention_mask=cross_attention_mask,
+            full_text_row_masked_out_mask=full_text_row_masked_out_mask,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            inputs_embeds=inputs_embeds,
+            **kwargs,
+        )
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=outputs.last_hidden_state,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+@auto_docstring(
+    custom_intro="""
+    The Mllama model which consists of a vision encoder and a language model.
+    """,
+)
+class MllamaForConditionalGeneration(MllamaPreTrainedModel, GenerationMixin):
+    # _tied_weights_keys = {"lm_head.weight": "model.language_moddel.embed_tokens.weight"}
+
+    def __init__(self, config: MllamaConfig):
+        super().__init__(config)
+        self.model = MllamaModel(config)
+        self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False)
+        self.post_init()
+
+    @can_return_tuple
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        pixel_values: torch.FloatTensor | None = None,
+        aspect_ratio_mask: torch.Tensor | None = None,
+        aspect_ratio_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        cross_attention_mask: torch.Tensor | None = None,
+        cross_attention_states: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        logits_to_keep: int | torch.Tensor = 0,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple | CausalLMOutputWithPast:
+        r"""
+        aspect_ratio_mask (`torch.Tensor` of shape `(batch_size, max_num_images, max_num_tiles)`, *optional*):
+            Mask to avoid performing attention on padding tiles. Mask values selected in `[0, 1]`:
+
+            - 1 for tiles that are **not masked**,
+            - 0 for tiles that are **masked**.
+        aspect_ratio_ids (`torch.Tensor` of shape `(batch_size, max_num_images)`, *optional*):
+            Aspect ratio ids used to select the appropriate precomputed tile embeddings based on the aspect ratio of each input image.
+            These ids correspond to indices in the model's list of supported aspect ratios, offset by 1.
+
+            For example, if the model supports aspect ratios [[1, 1], [1, 2], [2, 1]]:
+            - An image with aspect ratio [1, 1] would have ID 1
+            - An image with aspect ratio [1, 2] would have ID 2
+            - An image with aspect ratio [2, 1] would have ID 3
+
+            The id 0 is reserved for padding (i.e., no image).
+
+            If an image has aspect ratio [1, 2], that means it was split into 2 tiles horizontally, and its `aspect_ratio_id` would be 2.
+        cross_attention_mask (`torch.Tensor` of shape `(batch_size, seq_length, max_num_images, max_num_tiles)`, *optional*):
+            Cross-attention mask to control the interaction between text tokens and image tiles.
+            This 4D tensor defines which image tiles each text token should attend to.
+
+            For each text token (in seq_length):
+            - 1 indicates the token **should attend** to the corresponding image tile
+            - 0 indicates the token **should not attend** to the corresponding image tile
+        cross_attention_states (`torch.FloatTensor`, *optional*):
+            Output of the vision model, used for cross-attention. This tensor contains the processed image features that
+            the language model will attend to.
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+        Example:
+
+        ```python
+        >>> from PIL import Image
+        >>> import httpx
+        >>> from io import BytesIO
+        >>> from transformers import AutoProcessor, MllamaForConditionalGeneration
+
+        >>> checkpoint = "meta-llama/Llama-3.2-11B-Vision"
+        >>> model = MllamaForConditionalGeneration.from_pretrained(checkpoint)
+        >>> processor = AutoProcessor.from_pretrained(checkpoint)
+
+        >>> prompt = "<|image|>If I had to write a haiku for this one"
+        >>> url = "https://www.ilankelman.org/stopsigns/australia.jpg"
+        >>> with httpx.stream("GET", url) as response:
+        ...     image = Image.open(BytesIO(response.read()))
+
+        >>> inputs = processor(text=prompt, images=image, return_tensors="pt")
+
+        >>> # Generate
+        >>> output = model.generate(**inputs, max_new_tokens=15)
+
+        >>> prompt_len = inputs.input_ids.shape[-1]
+        >>> generated_ids = output[:, prompt_len:]
+        >>> generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)
+        >>> print(generated_text)
+        [', it would be:.\\nA stop sign in Chinatown.\\n']
+        ```
+        """
+        outputs = self.model(
+            input_ids=input_ids,
+            pixel_values=pixel_values,
+            aspect_ratio_mask=aspect_ratio_mask,
+            aspect_ratio_ids=aspect_ratio_ids,
+            cross_attention_mask=cross_attention_mask,
+            cross_attention_states=cross_attention_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            **kwargs,
+        )
+
+        hidden_states = outputs.last_hidden_state
+        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(logits, labels, self.config.text_config.vocab_size, **kwargs)
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids=None,
+        inputs_embeds=None,
+        attention_mask=None,
+        position_ids=None,
+        pixel_values=None,
+        aspect_ratio_ids=None,
+        aspect_ratio_mask=None,
+        cross_attention_mask=None,
+        past_key_values=None,
+        use_cache=False,
+        logits_to_keep=None,
+        is_first_iteration=False,
+        **kwargs,
+    ):
+        # Overwritten -- in specific circumstances we don't want to forward image inputs to the model
+
+        model_inputs = super().prepare_inputs_for_generation(
+            input_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            inputs_embeds=inputs_embeds,
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+            pixel_values=pixel_values,
+            aspect_ratio_ids=aspect_ratio_ids,
+            aspect_ratio_mask=aspect_ratio_mask,
+            cross_attention_mask=cross_attention_mask,
+            logits_to_keep=logits_to_keep,
+            is_first_iteration=is_first_iteration,
+            **kwargs,
+        )
+
+        # If we're in pre-fill or cacheless decoding step, then we need pixel_values and aspect ratios
+        # to compute image hidden states, otherwise they are cached within each cross attn layer
+        if not is_first_iteration and use_cache:
+            model_inputs["pixel_values"] = None
+            model_inputs["aspect_ratio_ids"] = None
+            model_inputs["aspect_ratio_mask"] = None
+
+        return model_inputs
+
+    def _update_model_kwargs_for_generation(self, outputs, model_kwargs, is_encoder_decoder, **kwargs):
+        cross_attention_mask_prev = model_kwargs.get("cross_attention_mask", None)
+        model_kwargs = super()._update_model_kwargs_for_generation(
+            outputs=outputs,
+            model_kwargs=model_kwargs,
+            is_encoder_decoder=is_encoder_decoder,
+            **kwargs,
+        )
+
+        # add cross-attn mask for new token
+        if cross_attention_mask_prev is not None:
+            model_kwargs["cross_attention_mask"] = torch.cat(
+                [cross_attention_mask_prev, cross_attention_mask_prev[:, -1:, ...]], dim=1
+            )
+        return model_kwargs
+
+
+__all__ = [
+    "MllamaForConditionalGeneration",
+    "MllamaForCausalLM",
+    "MllamaTextModel",
+    "MllamaVisionModel",
+    "MllamaPreTrainedModel",
+    "MllamaModel",
+]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/mllama/processing_mllama.py

@@ -0,0 +1,311 @@
+# Copyright 2024 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Processor class for Mllama."""
+
+import numpy as np
+
+from ...feature_extraction_utils import BatchFeature
+from ...image_utils import ImageInput, make_nested_list_of_images
+from ...processing_utils import ProcessingKwargs, ProcessorMixin, Unpack
+from ...tokenization_utils_base import PreTokenizedInput, TextInput
+from ...utils import auto_docstring
+
+
+class MllamaProcessorKwargs(ProcessingKwargs, total=False):
+    _defaults = {
+        "image_kwargs": {
+            "max_image_tiles": 4,
+        },
+    }
+
+
+def get_cross_attention_token_mask(input_ids: list[int], image_token_id: int) -> list[list[int]]:
+    """
+    Generate a cross-attention token mask for image tokens in the input sequence.
+
+    This function identifies the positions of image tokens in the input sequence and creates
+    a mask that defines which subsequent tokens each image token should attend to.
+
+    Args:
+        input_ids (list[int]): A list of token ids representing the input sequence.
+        image_token_id (int): The id of the token used to represent images in the sequence.
+
+    Returns:
+        list[list[int]]: A list of [start, end] pairs, where each pair represents the range
+        of tokens an image token should attend to.
+
+    Notes:
+        - If no image tokens are present, an empty list is returned.
+        - For a single image token, it attends to all subsequent tokens until the end of the sequence.
+        - For multiple image tokens, each attends to tokens up to the next image token or the end of the sequence.
+        - Consecutive image tokens are treated as a group and attend to all subsequent tokens together.
+    """
+
+    image_token_locations = [i for i, token in enumerate(input_ids) if token == image_token_id]
+
+    if len(image_token_locations) == 0:
+        return []
+
+    # only one image present, unmask until end of sequence
+    if len(image_token_locations) == 1:
+        return [[image_token_locations[0], -1]]
+
+    vision_masks = [[loc1, loc2] for loc1, loc2 in zip(image_token_locations[:-1], image_token_locations[1:])]
+
+    # last image will attend to all subsequent text
+    vision_masks.append([image_token_locations[-1], len(input_ids)])
+
+    # if there are two or more consecutive vision tokens,
+    # they should all attend to all subsequent
+    # text present
+    last_mask_end = vision_masks[-1][1]
+    for vision_mask in vision_masks[::-1]:
+        if vision_mask[0] == vision_mask[1] - 1:
+            vision_mask[1] = last_mask_end
+        last_mask_end = vision_mask[1]
+
+    return vision_masks
+
+
+def convert_sparse_cross_attention_mask_to_dense(
+    cross_attention_token_mask: list[list[list[int]]],
+    num_tiles: list[list[int]],
+    max_num_tiles: int,
+    length: int,
+) -> np.ndarray:
+    """
+    Convert the cross attention mask indices to a cross attention mask 4D array.
+
+    This function takes a sparse representation of cross attention masks and converts it to a dense 4D numpy array.
+    The sparse representation is a nested list structure that defines attention ranges for each image in each batch item.
+
+    Args:
+        cross_attention_token_mask (list[list[list[int]]]): A nested list structure where:
+            - The outer list represents the batch dimension.
+            - The middle list represents different images within each batch item.
+            - The inner list contains pairs of integers [start, end] representing token ranges for each image.
+        num_tiles (list[list[int]]): A nested list structure specifying the number of tiles for each image in each batch item.
+        max_num_tiles (int): The maximum possible number of tiles.
+        length (int): The total sequence length of the input.
+
+    Returns:
+        np.ndarray: A 4D numpy array of shape (batch_size, length, max_num_images, max_num_tiles)
+            The array contains `1` where attention is allowed and `0` where it is not.
+
+    Note:
+        - Special handling is done for cases where the end token is -1, which is interpreted as attending to the end of the sequence.
+    """
+
+    batch_size = len(cross_attention_token_mask)
+    max_num_images = max(len(masks) for masks in cross_attention_token_mask)
+
+    cross_attention_mask = np.zeros(
+        shape=(batch_size, length, max_num_images, max_num_tiles),
+        dtype=np.int64,
+    )
+
+    for sample_idx, (sample_masks, sample_num_tiles) in enumerate(zip(cross_attention_token_mask, num_tiles)):
+        for mask_idx, (locations, mask_num_tiles) in enumerate(zip(sample_masks, sample_num_tiles)):
+            if len(locations) == 2:
+                start, end = locations
+                end = min(end, length)
+                if end == -1:
+                    end = length
+                cross_attention_mask[sample_idx, start:end, mask_idx, :mask_num_tiles] = 1
+    return cross_attention_mask
+
+
+def build_string_from_input(prompt: str, bos_token: str, image_token: str) -> str:
+    """
+    Builds a string from the input prompt by adding `bos_token` if not already present.
+
+    Args:
+        prompt (`str`):
+            The input prompt string.
+        bos_token (`str`):
+            The beginning of sentence token to be added.
+        image_token (`str`):
+            The image token used to identify the start of an image sequence.
+
+    Returns:
+        str: The modified prompt string with the `bos_token` added if necessary.
+
+    Examples:
+        >>> build_string_from_input("Hello world", "<begin_of_text>", "<|image|>")
+        '<begin_of_text>Hello world'
+
+        >>> build_string_from_input("<|image|>Hello world", "<begin_of_text>", "<|image|>")
+        '<|image|><begin_of_text>Hello world'
+
+        >>> build_string_from_input("<begin_of_text>Hello world", "<begin_of_text>", "<|image|>")
+        '<begin_of_text>Hello world'
+    """
+
+    if bos_token in prompt:
+        return prompt
+
+    num_image_tokens_on_start = 0
+    while prompt.startswith(image_token):
+        prompt = prompt[len(image_token) :]
+        num_image_tokens_on_start += 1
+
+    return f"{image_token * num_image_tokens_on_start}{bos_token}{prompt}"
+
+
+@auto_docstring
+class MllamaProcessor(ProcessorMixin):
+    def __init__(self, image_processor, tokenizer, chat_template=None):
+        if not hasattr(tokenizer, "image_token"):
+            self.image_token = "<|image|>"
+            self.image_token_id = tokenizer.convert_tokens_to_ids(self.image_token)
+        else:
+            self.image_token = tokenizer.image_token
+            self.image_token_id = tokenizer.image_token_id
+
+        self.python_token = "<|python_tag|>"
+        self.python_token_id = tokenizer.convert_tokens_to_ids(self.python_token)
+        self.bos_token = tokenizer.bos_token
+        super().__init__(image_processor, tokenizer, chat_template=chat_template)
+
+    @auto_docstring
+    def __call__(
+        self,
+        images: ImageInput | None = None,
+        text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] | None = None,
+        **kwargs: Unpack[MllamaProcessorKwargs],
+    ) -> BatchFeature:
+        r"""
+        Returns:
+            [`BatchFeature`]: A [`BatchFeature`] with the following fields:
+
+            - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`.
+            - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when
+              `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not
+              `None`).
+            - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`.
+            TODO: add aspect_ratio_ids and aspect_ratio_mask and cross_attention_mask
+        """
+        if text is None and images is None:
+            raise ValueError("You must specify either text or images.")
+
+        output_kwargs = self._merge_kwargs(
+            MllamaProcessorKwargs,
+            tokenizer_init_kwargs=self.tokenizer.init_kwargs,
+            **kwargs,
+        )
+        return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None)
+
+        data = {}
+        if text is not None:
+            if isinstance(text, str):
+                text = [text]
+            elif not (isinstance(text, (list, tuple)) and all(isinstance(t, str) for t in text)):
+                raise ValueError("Invalid input text. Please provide a string, or a list of strings")
+            n_images_in_text = [t.count(self.image_token) for t in text]
+            text = [build_string_from_input(text_item, self.bos_token, self.image_token) for text_item in text]
+            encoding = self.tokenizer(text, **output_kwargs["text_kwargs"])
+            self._check_special_mm_tokens(text, encoding, modalities=["image"])
+            n_images_in_ids = [token_ids.count(self.image_token_id) for token_ids in encoding["input_ids"]]
+            data.update(encoding)
+
+        n_images_in_images = [0]
+        if images is not None:
+            images = self.image_processor.fetch_images(images)
+            images = make_nested_list_of_images(images)
+            n_images_in_images = [len(sample) for sample in images]
+
+        if text is not None:
+            if any(batch_img == 0 for batch_img in n_images_in_text) and not all(
+                batch_img == 0 for batch_img in n_images_in_text
+            ):
+                raise ValueError(
+                    "If a batch of text is provided, there should be either no images or at least one image per sample"
+                )
+            if sum(n_images_in_text) > 0 and (
+                n_images_in_images != n_images_in_text or n_images_in_ids != n_images_in_images
+            ):
+                if images is None:
+                    raise ValueError("No image were provided, but there are image tokens in the prompt")
+                else:
+                    add_message = ""
+                    if sum(n_images_in_images) == sum(n_images_in_text) and n_images_in_images != n_images_in_text:
+                        add_message = "Make sure to pass your images as a nested list, where each sub-list holds images per batch"
+                    elif n_images_in_ids != n_images_in_images:
+                        add_message = "If you activated truncation with `max_length`, increase the `max_length` so image tokens aren't cropped."
+
+                    raise ValueError(
+                        f"The number of image tokens in each text ({n_images_in_text}) should be the same as the "
+                        f"number of provided images per batch ({n_images_in_images}). {add_message}"
+                    )
+
+        if images is not None:
+            image_features = self.image_processor(images, **output_kwargs["images_kwargs"])
+            num_tiles = image_features.pop("num_tiles")
+            data.update(image_features)
+
+        # Create cross attention mask
+        if images is not None and text is not None:
+            cross_attention_token_mask = [
+                get_cross_attention_token_mask(token_ids, self.image_token_id) for token_ids in encoding["input_ids"]
+            ]
+            cross_attention_mask = convert_sparse_cross_attention_mask_to_dense(
+                cross_attention_token_mask,
+                num_tiles=num_tiles,
+                max_num_tiles=self.image_processor.max_image_tiles,
+                length=max(len(input_ids) for input_ids in encoding["input_ids"]),
+            )
+            data["cross_attention_mask"] = cross_attention_mask
+
+        return BatchFeature(data=data, tensor_type=return_tensors)
+
+    def post_process_image_text_to_text(
+        self, generated_outputs, skip_special_tokens=True, clean_up_tokenization_spaces=False, **kwargs
+    ):
+        """
+        Post-process the output of the model to decode the text.
+
+        Args:
+            generated_outputs (`torch.Tensor` or `np.ndarray`):
+                The output of the model `generate` function. The output is expected to be a tensor of shape `(batch_size, sequence_length)`
+                or `(sequence_length,)`.
+            skip_special_tokens (`bool`, *optional*, defaults to `True`):
+                Whether or not to remove special tokens in the output. Argument passed to the tokenizer's `batch_decode` method.
+            clean_up_tokenization_spaces (`bool`, *optional*, defaults to `False`):
+                Whether or not to clean up the tokenization spaces. Argument passed to the tokenizer's `batch_decode` method.
+            **kwargs:
+                Additional arguments to be passed to the tokenizer's `batch_decode method`.
+
+        Returns:
+            `list[str]`: The decoded text.
+        """
+        return self.tokenizer.batch_decode(
+            generated_outputs,
+            skip_special_tokens=skip_special_tokens,
+            clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+            **kwargs,
+        )
+
+    @property
+    def model_input_names(self):
+        tokenizer_input_names = self.tokenizer.model_input_names
+        image_processor_input_names = self.image_processor.model_input_names
+
+        # Remove `num_tiles`, it is popped and used only when processing. Make a copy of list when removing
+        # otherwise `self.image_processor.model_input_names` is also modified
+        image_processor_input_names = [name for name in image_processor_input_names if name != "num_tiles"]
+        return list(tokenizer_input_names + image_processor_input_names + ["cross_attention_mask"])
+
+
+__all__ = ["MllamaProcessor"]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/modernbert_decoder/__init__.py

@@ -0,0 +1,27 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+    from .configuration_modernbert_decoder import *
+    from .modeling_modernbert_decoder import *
+else:
+    import sys
+
+    _file = globals()["__file__"]
+    sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)

LTA_openwebtext_dualt/mini_owt_logdirichlet/logs/owt_t5_ultraclean10k_len1024_C4096_to_32768_exp_d768_l12_h12_gbs512_8gpu_40k_lr3e4_20260527_212316.log

@@ -0,0 +1,689 @@
+W0527 21:23:17.825000 10232 torch/distributed/run.py:792] 
+W0527 21:23:17.825000 10232 torch/distributed/run.py:792] *****************************************
+W0527 21:23:17.825000 10232 torch/distributed/run.py:792] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
+W0527 21:23:17.825000 10232 torch/distributed/run.py:792] *****************************************
+t-20260528052235-lvfnr-worker-0:10299:10299 [0] NCCL INFO NCCL_SOCKET_IFNAME set by environment to eth1
+t-20260528052235-lvfnr-worker-0:10299:10299 [0] NCCL INFO Bootstrap: Using eth1:10.82.64.8<0>
+t-20260528052235-lvfnr-worker-0:10299:10299 [0] NCCL INFO cudaDriverVersion 12080
+t-20260528052235-lvfnr-worker-0:10299:10299 [0] NCCL INFO NCCL version 2.25.1+cuda12.8
+t-20260528052235-lvfnr-worker-0:10299:10299 [0] NCCL INFO Comm config Blocking set to 1
+t-20260528052235-lvfnr-worker-0:10306:10306 [7] NCCL INFO cudaDriverVersion 12080
+t-20260528052235-lvfnr-worker-0:10306:10306 [7] NCCL INFO NCCL_SOCKET_IFNAME set by environment to eth1
+t-20260528052235-lvfnr-worker-0:10306:10306 [7] NCCL INFO Bootstrap: Using eth1:10.82.64.8<0>
+t-20260528052235-lvfnr-worker-0:10306:10306 [7] NCCL INFO NCCL version 2.25.1+cuda12.8
+t-20260528052235-lvfnr-worker-0:10302:10302 [3] NCCL INFO cudaDriverVersion 12080
+t-20260528052235-lvfnr-worker-0:10302:10302 [3] NCCL INFO NCCL_SOCKET_IFNAME set by environment to eth1
+t-20260528052235-lvfnr-worker-0:10306:10306 [7] NCCL INFO Comm config Blocking set to 1
+t-20260528052235-lvfnr-worker-0:10302:10302 [3] NCCL INFO Bootstrap: Using eth1:10.82.64.8<0>
+t-20260528052235-lvfnr-worker-0:10302:10302 [3] NCCL INFO NCCL version 2.25.1+cuda12.8
+t-20260528052235-lvfnr-worker-0:10302:10302 [3] NCCL INFO Comm config Blocking set to 1
+t-20260528052235-lvfnr-worker-0:10305:10305 [6] NCCL INFO cudaDriverVersion 12080
+t-20260528052235-lvfnr-worker-0:10305:10305 [6] NCCL INFO NCCL_SOCKET_IFNAME set by environment to eth1
+t-20260528052235-lvfnr-worker-0:10305:10305 [6] NCCL INFO Bootstrap: Using eth1:10.82.64.8<0>
+t-20260528052235-lvfnr-worker-0:10305:10305 [6] NCCL INFO NCCL version 2.25.1+cuda12.8
+t-20260528052235-lvfnr-worker-0:10305:10305 [6] NCCL INFO Comm config Blocking set to 1
+t-20260528052235-lvfnr-worker-0:10300:10300 [1] NCCL INFO cudaDriverVersion 12080
+t-20260528052235-lvfnr-worker-0:10300:10300 [1] NCCL INFO NCCL_SOCKET_IFNAME set by environment to eth1
+t-20260528052235-lvfnr-worker-0:10300:10300 [1] NCCL INFO Bootstrap: Using eth1:10.82.64.8<0>
+t-20260528052235-lvfnr-worker-0:10300:10300 [1] NCCL INFO NCCL version 2.25.1+cuda12.8
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO NET/Plugin: Loaded net plugin NCCL RDMA Plugin v9 (v9)
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO NET/Plugin: Loaded collnet plugin SHARP (v9)
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO Plugin Path : /opt/hpcx/nccl_rdma_sharp_plugin/lib/libnccl-net.so
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO P2P plugin v9 IBext_v9
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO NCCL_SOCKET_IFNAME set by environment to eth1
+t-20260528052235-lvfnr-worker-0:10300:10300 [1] NCCL INFO Comm config Blocking set to 1
+t-20260528052235-lvfnr-worker-0:10303:10303 [4] NCCL INFO cudaDriverVersion 12080
+t-20260528052235-lvfnr-worker-0:10303:10303 [4] NCCL INFO NCCL_SOCKET_IFNAME set by environment to eth1
+t-20260528052235-lvfnr-worker-0:10303:10303 [4] NCCL INFO Bootstrap: Using eth1:10.82.64.8<0>
+t-20260528052235-lvfnr-worker-0:10303:10303 [4] NCCL INFO NCCL version 2.25.1+cuda12.8
+t-20260528052235-lvfnr-worker-0:10303:10303 [4] NCCL INFO Comm config Blocking set to 1
+t-20260528052235-lvfnr-worker-0:10304:10304 [5] NCCL INFO cudaDriverVersion 12080
+t-20260528052235-lvfnr-worker-0:10304:10304 [5] NCCL INFO NCCL_SOCKET_IFNAME set by environment to eth1
+t-20260528052235-lvfnr-worker-0:10301:10301 [2] NCCL INFO cudaDriverVersion 12080
+t-20260528052235-lvfnr-worker-0:10301:10301 [2] NCCL INFO NCCL_SOCKET_IFNAME set by environment to eth1
+t-20260528052235-lvfnr-worker-0:10304:10304 [5] NCCL INFO Bootstrap: Using eth1:10.82.64.8<0>
+t-20260528052235-lvfnr-worker-0:10301:10301 [2] NCCL INFO Bootstrap: Using eth1:10.82.64.8<0>
+t-20260528052235-lvfnr-worker-0:10304:10304 [5] NCCL INFO NCCL version 2.25.1+cuda12.8
+t-20260528052235-lvfnr-worker-0:10301:10301 [2] NCCL INFO NCCL version 2.25.1+cuda12.8
+t-20260528052235-lvfnr-worker-0:10304:10304 [5] NCCL INFO Comm config Blocking set to 1
+t-20260528052235-lvfnr-worker-0:10301:10301 [2] NCCL INFO Comm config Blocking set to 1
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO NCCL_IB_PCI_RELAXED_ORDERING set by environment to 1.
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO NET/IB : Using [0]mlx5_1:1/RoCE [1]mlx5_4:1/RoCE [2]mlx5_5:1/RoCE [3]mlx5_6:1/RoCE [4]mlx5_7:1/RoCE [5]mlx5_8:1/RoCE [6]mlx5_9:1/RoCE [7]mlx5_10:1/RoCE [RO]; OOB eth1:10.82.64.8<0>
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO PROFILER/Plugin: Could not find: libnccl-profiler.so.
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO Using network IBext_v9
+t-20260528052235-lvfnr-worker-0:10306:10372 [7] NCCL INFO NET/Plugin: Loaded net plugin NCCL RDMA Plugin v9 (v9)
+t-20260528052235-lvfnr-worker-0:10306:10372 [7] NCCL INFO NET/Plugin: Loaded collnet plugin SHARP (v9)
+t-20260528052235-lvfnr-worker-0:10306:10372 [7] NCCL INFO Plugin Path : /opt/hpcx/nccl_rdma_sharp_plugin/lib/libnccl-net.so
+t-20260528052235-lvfnr-worker-0:10306:10372 [7] NCCL INFO P2P plugin v9 IBext_v9
+t-20260528052235-lvfnr-worker-0:10306:10372 [7] NCCL INFO NCCL_SOCKET_IFNAME set by environment to eth1
+t-20260528052235-lvfnr-worker-0:10302:10373 [3] NCCL INFO NET/Plugin: Loaded net plugin NCCL RDMA Plugin v9 (v9)
+t-20260528052235-lvfnr-worker-0:10302:10373 [3] NCCL INFO NET/Plugin: Loaded collnet plugin SHARP (v9)
+t-20260528052235-lvfnr-worker-0:10302:10373 [3] NCCL INFO Plugin Path : /opt/hpcx/nccl_rdma_sharp_plugin/lib/libnccl-net.so
+t-20260528052235-lvfnr-worker-0:10302:10373 [3] NCCL INFO P2P plugin v9 IBext_v9
+t-20260528052235-lvfnr-worker-0:10302:10373 [3] NCCL INFO NCCL_SOCKET_IFNAME set by environment to eth1
+t-20260528052235-lvfnr-worker-0:10305:10374 [6] NCCL INFO NET/Plugin: Loaded net plugin NCCL RDMA Plugin v9 (v9)
+t-20260528052235-lvfnr-worker-0:10305:10374 [6] NCCL INFO NET/Plugin: Loaded collnet plugin SHARP (v9)
+t-20260528052235-lvfnr-worker-0:10305:10374 [6] NCCL INFO Plugin Path : /opt/hpcx/nccl_rdma_sharp_plugin/lib/libnccl-net.so
+t-20260528052235-lvfnr-worker-0:10305:10374 [6] NCCL INFO P2P plugin v9 IBext_v9
+t-20260528052235-lvfnr-worker-0:10305:10374 [6] NCCL INFO NCCL_SOCKET_IFNAME set by environment to eth1
+t-20260528052235-lvfnr-worker-0:10306:10372 [7] NCCL INFO NCCL_IB_PCI_RELAXED_ORDERING set by environment to 1.
+t-20260528052235-lvfnr-worker-0:10306:10372 [7] NCCL INFO NET/IB : Using [0]mlx5_1:1/RoCE [1]mlx5_4:1/RoCE [2]mlx5_5:1/RoCE [3]mlx5_6:1/RoCE [4]mlx5_7:1/RoCE [5]mlx5_8:1/RoCE [6]mlx5_9:1/RoCE [7]mlx5_10:1/RoCE [RO]; OOB eth1:10.82.64.8<0>
+t-20260528052235-lvfnr-worker-0:10306:10372 [7] NCCL INFO PROFILER/Plugin: Could not find: libnccl-profiler.so.
+t-20260528052235-lvfnr-worker-0:10306:10372 [7] NCCL INFO Using network IBext_v9
+t-20260528052235-lvfnr-worker-0:10302:10373 [3] NCCL INFO NCCL_IB_PCI_RELAXED_ORDERING set by environment to 1.
+t-20260528052235-lvfnr-worker-0:10302:10373 [3] NCCL INFO NET/IB : Using [0]mlx5_1:1/RoCE [1]mlx5_4:1/RoCE [2]mlx5_5:1/RoCE [3]mlx5_6:1/RoCE [4]mlx5_7:1/RoCE [5]mlx5_8:1/RoCE [6]mlx5_9:1/RoCE [7]mlx5_10:1/RoCE [RO]; OOB eth1:10.82.64.8<0>
+t-20260528052235-lvfnr-worker-0:10302:10373 [3] NCCL INFO PROFILER/Plugin: Could not find: libnccl-profiler.so.
+t-20260528052235-lvfnr-worker-0:10302:10373 [3] NCCL INFO Using network IBext_v9
+t-20260528052235-lvfnr-worker-0:10305:10374 [6] NCCL INFO NCCL_IB_PCI_RELAXED_ORDERING set by environment to 1.
+t-20260528052235-lvfnr-worker-0:10305:10374 [6] NCCL INFO NET/IB : Using [0]mlx5_1:1/RoCE [1]mlx5_4:1/RoCE [2]mlx5_5:1/RoCE [3]mlx5_6:1/RoCE [4]mlx5_7:1/RoCE [5]mlx5_8:1/RoCE [6]mlx5_9:1/RoCE [7]mlx5_10:1/RoCE [RO]; OOB eth1:10.82.64.8<0>
+t-20260528052235-lvfnr-worker-0:10305:10374 [6] NCCL INFO PROFILER/Plugin: Could not find: libnccl-profiler.so.
+t-20260528052235-lvfnr-worker-0:10305:10374 [6] NCCL INFO Using network IBext_v9
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO ncclCommInitRankConfig comm 0x95c7e40 rank 0 nranks 8 cudaDev 0 nvmlDev 0 busId 65040 commId 0x1a57b0265de3f8e1 - Init START
+t-20260528052235-lvfnr-worker-0:10303:10377 [4] NCCL INFO NET/Plugin: Loaded net plugin NCCL RDMA Plugin v9 (v9)
+t-20260528052235-lvfnr-worker-0:10303:10377 [4] NCCL INFO NET/Plugin: Loaded collnet plugin SHARP (v9)
+t-20260528052235-lvfnr-worker-0:10303:10377 [4] NCCL INFO Plugin Path : /opt/hpcx/nccl_rdma_sharp_plugin/lib/libnccl-net.so
+t-20260528052235-lvfnr-worker-0:10303:10377 [4] NCCL INFO P2P plugin v9 IBext_v9
+t-20260528052235-lvfnr-worker-0:10303:10377 [4] NCCL INFO NCCL_SOCKET_IFNAME set by environment to eth1
+t-20260528052235-lvfnr-worker-0:10300:10375 [1] NCCL INFO NET/Plugin: Loaded net plugin NCCL RDMA Plugin v9 (v9)
+t-20260528052235-lvfnr-worker-0:10300:10375 [1] NCCL INFO NET/Plugin: Loaded collnet plugin SHARP (v9)
+t-20260528052235-lvfnr-worker-0:10300:10375 [1] NCCL INFO Plugin Path : /opt/hpcx/nccl_rdma_sharp_plugin/lib/libnccl-net.so
+t-20260528052235-lvfnr-worker-0:10300:10375 [1] NCCL INFO P2P plugin v9 IBext_v9
+t-20260528052235-lvfnr-worker-0:10300:10375 [1] NCCL INFO NCCL_SOCKET_IFNAME set by environment to eth1
+t-20260528052235-lvfnr-worker-0:10304:10381 [5] NCCL INFO NET/Plugin: Loaded net plugin NCCL RDMA Plugin v9 (v9)
+t-20260528052235-lvfnr-worker-0:10304:10381 [5] NCCL INFO NET/Plugin: Loaded collnet plugin SHARP (v9)
+t-20260528052235-lvfnr-worker-0:10304:10381 [5] NCCL INFO Plugin Path : /opt/hpcx/nccl_rdma_sharp_plugin/lib/libnccl-net.so
+t-20260528052235-lvfnr-worker-0:10304:10381 [5] NCCL INFO P2P plugin v9 IBext_v9
+t-20260528052235-lvfnr-worker-0:10304:10381 [5] NCCL INFO NCCL_SOCKET_IFNAME set by environment to eth1
+t-20260528052235-lvfnr-worker-0:10301:10382 [2] NCCL INFO NET/Plugin: Loaded net plugin NCCL RDMA Plugin v9 (v9)
+t-20260528052235-lvfnr-worker-0:10301:10382 [2] NCCL INFO NET/Plugin: Loaded collnet plugin SHARP (v9)
+t-20260528052235-lvfnr-worker-0:10301:10382 [2] NCCL INFO Plugin Path : /opt/hpcx/nccl_rdma_sharp_plugin/lib/libnccl-net.so
+t-20260528052235-lvfnr-worker-0:10301:10382 [2] NCCL INFO P2P plugin v9 IBext_v9
+t-20260528052235-lvfnr-worker-0:10301:10382 [2] NCCL INFO NCCL_SOCKET_IFNAME set by environment to eth1
+t-20260528052235-lvfnr-worker-0:10306:10372 [7] NCCL INFO ncclCommInitRankConfig comm 0xa010c30 rank 7 nranks 8 cudaDev 7 nvmlDev 7 busId 75020 commId 0x1a57b0265de3f8e1 - Init START
+t-20260528052235-lvfnr-worker-0:10302:10373 [3] NCCL INFO ncclCommInitRankConfig comm 0x8d69240 rank 3 nranks 8 cudaDev 3 nvmlDev 3 busId 6b020 commId 0x1a57b0265de3f8e1 - Init START
+t-20260528052235-lvfnr-worker-0:10305:10374 [6] NCCL INFO ncclCommInitRankConfig comm 0x919e3e0 rank 6 nranks 8 cudaDev 6 nvmlDev 6 busId 73020 commId 0x1a57b0265de3f8e1 - Init START
+t-20260528052235-lvfnr-worker-0:10306:10372 [7] NCCL INFO RAS client listening socket at ::1<28028>
+t-20260528052235-lvfnr-worker-0:10300:10375 [1] NCCL INFO NCCL_IB_PCI_RELAXED_ORDERING set by environment to 1.
+t-20260528052235-lvfnr-worker-0:10300:10375 [1] NCCL INFO NET/IB : Using [0]mlx5_1:1/RoCE [1]mlx5_4:1/RoCE [2]mlx5_5:1/RoCE [3]mlx5_6:1/RoCE [4]mlx5_7:1/RoCE [5]mlx5_8:1/RoCE [6]mlx5_9:1/RoCE [7]mlx5_10:1/RoCE [RO]; OOB eth1:10.82.64.8<0>
+t-20260528052235-lvfnr-worker-0:10304:10381 [5] NCCL INFO NCCL_IB_PCI_RELAXED_ORDERING set by environment to 1.
+t-20260528052235-lvfnr-worker-0:10304:10381 [5] NCCL INFO NET/IB : Using [0]mlx5_1:1/RoCE [1]mlx5_4:1/RoCE [2]mlx5_5:1/RoCE [3]mlx5_6:1/RoCE [4]mlx5_7:1/RoCE [5]mlx5_8:1/RoCE [6]mlx5_9:1/RoCE [7]mlx5_10:1/RoCE [RO]; OOB eth1:10.82.64.8<0>
+t-20260528052235-lvfnr-worker-0:10301:10382 [2] NCCL INFO NCCL_IB_PCI_RELAXED_ORDERING set by environment to 1.
+t-20260528052235-lvfnr-worker-0:10301:10382 [2] NCCL INFO NET/IB : Using [0]mlx5_1:1/RoCE [1]mlx5_4:1/RoCE [2]mlx5_5:1/RoCE [3]mlx5_6:1/RoCE [4]mlx5_7:1/RoCE [5]mlx5_8:1/RoCE [6]mlx5_9:1/RoCE [7]mlx5_10:1/RoCE [RO]; OOB eth1:10.82.64.8<0>
+t-20260528052235-lvfnr-worker-0:10303:10377 [4] NCCL INFO NCCL_IB_PCI_RELAXED_ORDERING set by environment to 1.
+t-20260528052235-lvfnr-worker-0:10303:10377 [4] NCCL INFO NET/IB : Using [0]mlx5_1:1/RoCE [1]mlx5_4:1/RoCE [2]mlx5_5:1/RoCE [3]mlx5_6:1/RoCE [4]mlx5_7:1/RoCE [5]mlx5_8:1/RoCE [6]mlx5_9:1/RoCE [7]mlx5_10:1/RoCE [RO]; OOB eth1:10.82.64.8<0>
+t-20260528052235-lvfnr-worker-0:10300:10375 [1] NCCL INFO PROFILER/Plugin: Could not find: libnccl-profiler.so.
+t-20260528052235-lvfnr-worker-0:10300:10375 [1] NCCL INFO Using network IBext_v9
+t-20260528052235-lvfnr-worker-0:10304:10381 [5] NCCL INFO PROFILER/Plugin: Could not find: libnccl-profiler.so.
+t-20260528052235-lvfnr-worker-0:10304:10381 [5] NCCL INFO Using network IBext_v9
+t-20260528052235-lvfnr-worker-0:10301:10382 [2] NCCL INFO PROFILER/Plugin: Could not find: libnccl-profiler.so.
+t-20260528052235-lvfnr-worker-0:10301:10382 [2] NCCL INFO Using network IBext_v9
+t-20260528052235-lvfnr-worker-0:10303:10377 [4] NCCL INFO PROFILER/Plugin: Could not find: libnccl-profiler.so.
+t-20260528052235-lvfnr-worker-0:10303:10377 [4] NCCL INFO Using network IBext_v9
+t-20260528052235-lvfnr-worker-0:10300:10375 [1] NCCL INFO ncclCommInitRankConfig comm 0xcebc1b0 rank 1 nranks 8 cudaDev 1 nvmlDev 1 busId 67020 commId 0x1a57b0265de3f8e1 - Init START
+t-20260528052235-lvfnr-worker-0:10304:10381 [5] NCCL INFO ncclCommInitRankConfig comm 0xce66c40 rank 5 nranks 8 cudaDev 5 nvmlDev 5 busId 71020 commId 0x1a57b0265de3f8e1 - Init START
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO RAS client listening socket at ::1<28028>
+t-20260528052235-lvfnr-worker-0:10305:10374 [6] NCCL INFO RAS client listening socket at ::1<28028>
+t-20260528052235-lvfnr-worker-0:10301:10382 [2] NCCL INFO ncclCommInitRankConfig comm 0xcec7380 rank 2 nranks 8 cudaDev 2 nvmlDev 2 busId 69020 commId 0x1a57b0265de3f8e1 - Init START
+t-20260528052235-lvfnr-worker-0:10300:10375 [1] NCCL INFO RAS client listening socket at ::1<28028>
+t-20260528052235-lvfnr-worker-0:10301:10382 [2] NCCL INFO RAS client listening socket at ::1<28028>
+t-20260528052235-lvfnr-worker-0:10303:10377 [4] NCCL INFO ncclCommInitRankConfig comm 0xc7a6880 rank 4 nranks 8 cudaDev 4 nvmlDev 4 busId 6f020 commId 0x1a57b0265de3f8e1 - Init START
+t-20260528052235-lvfnr-worker-0:10304:10381 [5] NCCL INFO RAS client listening socket at ::1<28028>
+t-20260528052235-lvfnr-worker-0:10302:10373 [3] NCCL INFO RAS client listening socket at ::1<28028>
+t-20260528052235-lvfnr-worker-0:10303:10377 [4] NCCL INFO RAS client listening socket at ::1<28028>
+t-20260528052235-lvfnr-worker-0:10301:10382 [2] NCCL INFO Bootstrap timings total 0.001228 (create 0.000020, send 0.000066, recv 0.000082, ring 0.000786, delay 0.000001)
+t-20260528052235-lvfnr-worker-0:10302:10373 [3] NCCL INFO Bootstrap timings total 0.617195 (create 0.000025, send 0.000072, recv 0.616667, ring 0.000100, delay 0.000001)
+t-20260528052235-lvfnr-worker-0:10300:10375 [1] NCCL INFO Bootstrap timings total 0.002286 (create 0.000024, send 0.000068, recv 0.001100, ring 0.000791, delay 0.000001)
+t-20260528052235-lvfnr-worker-0:10303:10377 [4] NCCL INFO Bootstrap timings total 0.000677 (create 0.000022, send 0.000070, recv 0.000166, ring 0.000115, delay 0.000001)
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO Bootstrap timings total 0.900655 (create 0.000023, send 0.000079, recv 0.898502, ring 0.001665, delay 0.000001)
+t-20260528052235-lvfnr-worker-0:10304:10381 [5] NCCL INFO Bootstrap timings total 0.001774 (create 0.000021, send 0.000065, recv 0.000052, ring 0.000117, delay 0.000001)
+t-20260528052235-lvfnr-worker-0:10305:10374 [6] NCCL INFO Bootstrap timings total 0.604435 (create 0.000021, send 0.000065, recv 0.000060, ring 0.001248, delay 0.000001)
+t-20260528052235-lvfnr-worker-0:10306:10372 [7] NCCL INFO Bootstrap timings total 0.628573 (create 0.000023, send 0.000066, recv 0.000105, ring 0.554812, delay 0.000000)
+t-20260528052235-lvfnr-worker-0:10301:10382 [2] NCCL INFO MNNVL busId 0x69020 fabric UUID 0.0 cliqueId 0x0 state 3 healthMask 0x0
+t-20260528052235-lvfnr-worker-0:10300:10375 [1] NCCL INFO MNNVL busId 0x67020 fabric UUID 0.0 cliqueId 0x0 state 3 healthMask 0x0
+t-20260528052235-lvfnr-worker-0:10302:10373 [3] NCCL INFO MNNVL busId 0x6b020 fabric UUID 0.0 cliqueId 0x0 state 3 healthMask 0x0
+t-20260528052235-lvfnr-worker-0:10304:10381 [5] NCCL INFO MNNVL busId 0x71020 fabric UUID 0.0 cliqueId 0x0 state 3 healthMask 0x0
+t-20260528052235-lvfnr-worker-0:10303:10377 [4] NCCL INFO MNNVL busId 0x6f020 fabric UUID 0.0 cliqueId 0x0 state 3 healthMask 0x0
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO MNNVL busId 0x65040 fabric UUID 0.0 cliqueId 0x0 state 3 healthMask 0x0
+t-20260528052235-lvfnr-worker-0:10305:10374 [6] NCCL INFO MNNVL busId 0x73020 fabric UUID 0.0 cliqueId 0x0 state 3 healthMask 0x0
+t-20260528052235-lvfnr-worker-0:10306:10372 [7] NCCL INFO MNNVL busId 0x75020 fabric UUID 0.0 cliqueId 0x0 state 3 healthMask 0x0
+t-20260528052235-lvfnr-worker-0:10301:10382 [2] NCCL INFO NCCL_TOPO_FILE set by environment to /var/run/nvidia-topologyd/virtualTopology.xml
+t-20260528052235-lvfnr-worker-0:10306:10372 [7] NCCL INFO NCCL_TOPO_FILE set by environment to /var/run/nvidia-topologyd/virtualTopology.xml
+t-20260528052235-lvfnr-worker-0:10304:10381 [5] NCCL INFO NCCL_TOPO_FILE set by environment to /var/run/nvidia-topologyd/virtualTopology.xml
+t-20260528052235-lvfnr-worker-0:10303:10377 [4] NCCL INFO NCCL_TOPO_FILE set by environment to /var/run/nvidia-topologyd/virtualTopology.xml
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO NCCL_TOPO_FILE set by environment to /var/run/nvidia-topologyd/virtualTopology.xml
+t-20260528052235-lvfnr-worker-0:10300:10375 [1] NCCL INFO NCCL_TOPO_FILE set by environment to /var/run/nvidia-topologyd/virtualTopology.xml
+t-20260528052235-lvfnr-worker-0:10302:10373 [3] NCCL INFO NCCL_TOPO_FILE set by environment to /var/run/nvidia-topologyd/virtualTopology.xml
+t-20260528052235-lvfnr-worker-0:10305:10374 [6] NCCL INFO NCCL_TOPO_FILE set by environment to /var/run/nvidia-topologyd/virtualTopology.xml
+t-20260528052235-lvfnr-worker-0:10301:10382 [2] NCCL INFO Setting affinity for GPU 2 to 03ffffff,ffffffff,ffffffff
+t-20260528052235-lvfnr-worker-0:10301:10382 [2] NCCL INFO NVLS multicast support is available on dev 2
+t-20260528052235-lvfnr-worker-0:10304:10381 [5] NCCL INFO Setting affinity for GPU 5 to 0fffff,ffffffff,ffffffff,fc000000,00000000,00000000
+t-20260528052235-lvfnr-worker-0:10300:10375 [1] NCCL INFO Setting affinity for GPU 1 to 03ffffff,ffffffff,ffffffff
+t-20260528052235-lvfnr-worker-0:10302:10373 [3] NCCL INFO Setting affinity for GPU 3 to 03ffffff,ffffffff,ffffffff
+t-20260528052235-lvfnr-worker-0:10300:10375 [1] NCCL INFO NVLS multicast support is available on dev 1
+t-20260528052235-lvfnr-worker-0:10302:10373 [3] NCCL INFO NVLS multicast support is available on dev 3
+t-20260528052235-lvfnr-worker-0:10304:10381 [5] NCCL INFO NVLS multicast support is available on dev 5
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO Setting affinity for GPU 0 to 03ffffff,ffffffff,ffffffff
+t-20260528052235-lvfnr-worker-0:10303:10377 [4] NCCL INFO Setting affinity for GPU 4 to 0fffff,ffffffff,ffffffff,fc000000,00000000,00000000
+t-20260528052235-lvfnr-worker-0:10306:10372 [7] NCCL INFO Setting affinity for GPU 7 to 0fffff,ffffffff,ffffffff,fc000000,00000000,00000000
+t-20260528052235-lvfnr-worker-0:10305:10374 [6] NCCL INFO Setting affinity for GPU 6 to 0fffff,ffffffff,ffffffff,fc000000,00000000,00000000
+t-20260528052235-lvfnr-worker-0:10303:10377 [4] NCCL INFO NVLS multicast support is available on dev 4
+t-20260528052235-lvfnr-worker-0:10306:10372 [7] NCCL INFO NVLS multicast support is available on dev 7
+t-20260528052235-lvfnr-worker-0:10305:10374 [6] NCCL INFO NVLS multicast support is available on dev 6
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO NVLS multicast support is available on dev 0
+t-20260528052235-lvfnr-worker-0:10306:10372 [7] NCCL INFO comm 0xa010c30 rank 7 nRanks 8 nNodes 1 localRanks 8 localRank 7 MNNVL 0
+t-20260528052235-lvfnr-worker-0:10305:10374 [6] NCCL INFO comm 0x919e3e0 rank 6 nRanks 8 nNodes 1 localRanks 8 localRank 6 MNNVL 0
+t-20260528052235-lvfnr-worker-0:10300:10375 [1] NCCL INFO comm 0xcebc1b0 rank 1 nRanks 8 nNodes 1 localRanks 8 localRank 1 MNNVL 0
+t-20260528052235-lvfnr-worker-0:10304:10381 [5] NCCL INFO comm 0xce66c40 rank 5 nRanks 8 nNodes 1 localRanks 8 localRank 5 MNNVL 0
+t-20260528052235-lvfnr-worker-0:10301:10382 [2] NCCL INFO comm 0xcec7380 rank 2 nRanks 8 nNodes 1 localRanks 8 localRank 2 MNNVL 0
+t-20260528052235-lvfnr-worker-0:10303:10377 [4] NCCL INFO comm 0xc7a6880 rank 4 nRanks 8 nNodes 1 localRanks 8 localRank 4 MNNVL 0
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO comm 0x95c7e40 rank 0 nRanks 8 nNodes 1 localRanks 8 localRank 0 MNNVL 0
+t-20260528052235-lvfnr-worker-0:10302:10373 [3] NCCL INFO comm 0x8d69240 rank 3 nRanks 8 nNodes 1 localRanks 8 localRank 3 MNNVL 0
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO Channel 00/24 : 0 1 2 3 4 5 6 7
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO Channel 01/24 : 0 1 2 3 4 5 6 7
+t-20260528052235-lvfnr-worker-0:10303:10377 [4] NCCL INFO Trees [0] 5/-1/-1->4->3 [1] 5/-1/-1->4->3 [2] 5/-1/-1->4->3 [3] 5/-1/-1->4->3 [4] 5/-1/-1->4->3 [5] 5/-1/-1->4->3 [6] 5/-1/-1->4->3 [7] 5/-1/-1->4->3 [8] 5/-1/-1->4->3 [9] 5/-1/-1->4->3 [10] 5/-1/-1->4->3 [11] 5/-1/-1->4->3 [12] 5/-1/-1->4->3 [13] 5/-1/-1->4->3 [14] 5/-1/-1->4->3 [15] 5/-1/-1->4->3 [16] 5/-1/-1->4->3 [17] 5/-1/-1->4->3 [18] 5/-1/-1->4->3 [19] 5/-1/-1->4->3 [20] 5/-1/-1->4->3 [21] 5/-1/-1->4->3 [22] 5/-1/-1->4->3 [23] 5/-1/-1->4->3
+t-20260528052235-lvfnr-worker-0:10306:10372 [7] NCCL INFO Trees [0] -1/-1/-1->7->6 [1] -1/-1/-1->7->6 [2] -1/-1/-1->7->6 [3] -1/-1/-1->7->6 [4] -1/-1/-1->7->6 [5] -1/-1/-1->7->6 [6] -1/-1/-1->7->6 [7] -1/-1/-1->7->6 [8] -1/-1/-1->7->6 [9] -1/-1/-1->7->6 [10] -1/-1/-1->7->6 [11] -1/-1/-1->7->6 [12] -1/-1/-1->7->6 [13] -1/-1/-1->7->6 [14] -1/-1/-1->7->6 [15] -1/-1/-1->7->6 [16] -1/-1/-1->7->6 [17] -1/-1/-1->7->6 [18] -1/-1/-1->7->6 [19] -1/-1/-1->7->6 [20] -1/-1/-1->7->6 [21] -1/-1/-1->7->6 [22] -1/-1/-1->7->6 [23] -1/-1/-1->7->6
+t-20260528052235-lvfnr-worker-0:10302:10373 [3] NCCL INFO Trees [0] 4/-1/-1->3->2 [1] 4/-1/-1->3->2 [2] 4/-1/-1->3->2 [3] 4/-1/-1->3->2 [4] 4/-1/-1->3->2 [5] 4/-1/-1->3->2 [6] 4/-1/-1->3->2 [7] 4/-1/-1->3->2 [8] 4/-1/-1->3->2 [9] 4/-1/-1->3->2 [10] 4/-1/-1->3->2 [11] 4/-1/-1->3->2 [12] 4/-1/-1->3->2 [13] 4/-1/-1->3->2 [14] 4/-1/-1->3->2 [15] 4/-1/-1->3->2 [16] 4/-1/-1->3->2 [17] 4/-1/-1->3->2 [18] 4/-1/-1->3->2 [19] 4/-1/-1->3->2 [20] 4/-1/-1->3->2 [21] 4/-1/-1->3->2 [22] 4/-1/-1->3->2 [23] 4/-1/-1->3->2
+t-20260528052235-lvfnr-worker-0:10300:10375 [1] NCCL INFO Trees [0] 2/-1/-1->1->0 [1] 2/-1/-1->1->0 [2] 2/-1/-1->1->0 [3] 2/-1/-1->1->0 [4] 2/-1/-1->1->0 [5] 2/-1/-1->1->0 [6] 2/-1/-1->1->0 [7] 2/-1/-1->1->0 [8] 2/-1/-1->1->0 [9] 2/-1/-1->1->0 [10] 2/-1/-1->1->0 [11] 2/-1/-1->1->0 [12] 2/-1/-1->1->0 [13] 2/-1/-1->1->0 [14] 2/-1/-1->1->0 [15] 2/-1/-1->1->0 [16] 2/-1/-1->1->0 [17] 2/-1/-1->1->0 [18] 2/-1/-1->1->0 [19] 2/-1/-1->1->0 [20] 2/-1/-1->1->0 [21] 2/-1/-1->1->0 [22] 2/-1/-1->1->0 [23] 2/-1/-1->1->0
+t-20260528052235-lvfnr-worker-0:10305:10374 [6] NCCL INFO Trees [0] 7/-1/-1->6->5 [1] 7/-1/-1->6->5 [2] 7/-1/-1->6->5 [3] 7/-1/-1->6->5 [4] 7/-1/-1->6->5 [5] 7/-1/-1->6->5 [6] 7/-1/-1->6->5 [7] 7/-1/-1->6->5 [8] 7/-1/-1->6->5 [9] 7/-1/-1->6->5 [10] 7/-1/-1->6->5 [11] 7/-1/-1->6->5 [12] 7/-1/-1->6->5 [13] 7/-1/-1->6->5 [14] 7/-1/-1->6->5 [15] 7/-1/-1->6->5 [16] 7/-1/-1->6->5 [17] 7/-1/-1->6->5 [18] 7/-1/-1->6->5 [19] 7/-1/-1->6->5 [20] 7/-1/-1->6->5 [21] 7/-1/-1->6->5 [22] 7/-1/-1->6->5 [23] 7/-1/-1->6->5
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO Channel 02/24 : 0 1 2 3 4 5 6 7
+t-20260528052235-lvfnr-worker-0:10303:10377 [4] NCCL INFO P2P Chunksize set to 524288
+t-20260528052235-lvfnr-worker-0:10306:10372 [7] NCCL INFO P2P Chunksize set to 524288
+t-20260528052235-lvfnr-worker-0:10302:10373 [3] NCCL INFO P2P Chunksize set to 524288
+t-20260528052235-lvfnr-worker-0:10304:10381 [5] NCCL INFO Trees [0] 6/-1/-1->5->4 [1] 6/-1/-1->5->4 [2] 6/-1/-1->5->4 [3] 6/-1/-1->5->4 [4] 6/-1/-1->5->4 [5] 6/-1/-1->5->4 [6] 6/-1/-1->5->4 [7] 6/-1/-1->5->4 [8] 6/-1/-1->5->4 [9] 6/-1/-1->5->4 [10] 6/-1/-1->5->4 [11] 6/-1/-1->5->4 [12] 6/-1/-1->5->4 [13] 6/-1/-1->5->4 [14] 6/-1/-1->5->4 [15] 6/-1/-1->5->4 [16] 6/-1/-1->5->4 [17] 6/-1/-1->5->4 [18] 6/-1/-1->5->4 [19] 6/-1/-1->5->4 [20] 6/-1/-1->5->4 [21] 6/-1/-1->5->4 [22] 6/-1/-1->5->4 [23] 6/-1/-1->5->4
+t-20260528052235-lvfnr-worker-0:10300:10375 [1] NCCL INFO P2P Chunksize set to 524288
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO Channel 03/24 : 0 1 2 3 4 5 6 7
+t-20260528052235-lvfnr-worker-0:10305:10374 [6] NCCL INFO P2P Chunksize set to 524288
+t-20260528052235-lvfnr-worker-0:10301:10382 [2] NCCL INFO Trees [0] 3/-1/-1->2->1 [1] 3/-1/-1->2->1 [2] 3/-1/-1->2->1 [3] 3/-1/-1->2->1 [4] 3/-1/-1->2->1 [5] 3/-1/-1->2->1 [6] 3/-1/-1->2->1 [7] 3/-1/-1->2->1 [8] 3/-1/-1->2->1 [9] 3/-1/-1->2->1 [10] 3/-1/-1->2->1 [11] 3/-1/-1->2->1 [12] 3/-1/-1->2->1 [13] 3/-1/-1->2->1 [14] 3/-1/-1->2->1 [15] 3/-1/-1->2->1 [16] 3/-1/-1->2->1 [17] 3/-1/-1->2->1 [18] 3/-1/-1->2->1 [19] 3/-1/-1->2->1 [20] 3/-1/-1->2->1 [21] 3/-1/-1->2->1 [22] 3/-1/-1->2->1 [23] 3/-1/-1->2->1
+t-20260528052235-lvfnr-worker-0:10304:10381 [5] NCCL INFO P2P Chunksize set to 524288
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO Channel 04/24 : 0 1 2 3 4 5 6 7
+t-20260528052235-lvfnr-worker-0:10301:10382 [2] NCCL INFO P2P Chunksize set to 524288
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO Channel 05/24 : 0 1 2 3 4 5 6 7
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO Channel 06/24 : 0 1 2 3 4 5 6 7
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO Channel 07/24 : 0 1 2 3 4 5 6 7
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO Channel 08/24 : 0 1 2 3 4 5 6 7
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO Channel 09/24 : 0 1 2 3 4 5 6 7
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO Channel 10/24 : 0 1 2 3 4 5 6 7
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO Channel 11/24 : 0 1 2 3 4 5 6 7
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO Channel 12/24 : 0 1 2 3 4 5 6 7
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO Channel 13/24 : 0 1 2 3 4 5 6 7
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO Channel 14/24 : 0 1 2 3 4 5 6 7
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO Channel 15/24 : 0 1 2 3 4 5 6 7
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO Channel 16/24 : 0 1 2 3 4 5 6 7
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO Channel 17/24 : 0 1 2 3 4 5 6 7
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO Channel 18/24 : 0 1 2 3 4 5 6 7
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO Channel 19/24 : 0 1 2 3 4 5 6 7
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO Channel 20/24 : 0 1 2 3 4 5 6 7
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO Channel 21/24 : 0 1 2 3 4 5 6 7
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO Channel 22/24 : 0 1 2 3 4 5 6 7
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO Channel 23/24 : 0 1 2 3 4 5 6 7
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO Trees [0] 1/-1/-1->0->-1 [1] 1/-1/-1->0->-1 [2] 1/-1/-1->0->-1 [3] 1/-1/-1->0->-1 [4] 1/-1/-1->0->-1 [5] 1/-1/-1->0->-1 [6] 1/-1/-1->0->-1 [7] 1/-1/-1->0->-1 [8] 1/-1/-1->0->-1 [9] 1/-1/-1->0->-1 [10] 1/-1/-1->0->-1 [11] 1/-1/-1->0->-1 [12] 1/-1/-1->0->-1 [13] 1/-1/-1->0->-1 [14] 1/-1/-1->0->-1 [15] 1/-1/-1->0->-1 [16] 1/-1/-1->0->-1 [17] 1/-1/-1->0->-1 [18] 1/-1/-1->0->-1 [19] 1/-1/-1->0->-1 [20] 1/-1/-1->0->-1 [21] 1/-1/-1->0->-1 [22] 1/-1/-1->0->-1 [23] 1/-1/-1->0->-1
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO P2P Chunksize set to 524288
+t-20260528052235-lvfnr-worker-0:10303:10451 [4] NCCL INFO [Proxy Service] Device 4 CPU core 96
+t-20260528052235-lvfnr-worker-0:10303:10452 [4] NCCL INFO [Proxy Service UDS] Device 4 CPU core 98
+t-20260528052235-lvfnr-worker-0:10300:10453 [1] NCCL INFO [Proxy Service] Device 1 CPU core 60
+t-20260528052235-lvfnr-worker-0:10300:10454 [1] NCCL INFO [Proxy Service UDS] Device 1 CPU core 62
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO Check P2P Type intraNodeP2pSupport 1 directMode 0
+t-20260528052235-lvfnr-worker-0:10299:10455 [0] NCCL INFO [Proxy Service] Device 0 CPU core 2
+t-20260528052235-lvfnr-worker-0:10299:10456 [0] NCCL INFO [Proxy Service UDS] Device 0 CPU core 4
+t-20260528052235-lvfnr-worker-0:10301:10457 [2] NCCL INFO [Proxy Service] Device 2 CPU core 70
+t-20260528052235-lvfnr-worker-0:10301:10458 [2] NCCL INFO [Proxy Service UDS] Device 2 CPU core 72
+t-20260528052235-lvfnr-worker-0:10306:10459 [7] NCCL INFO [Proxy Service] Device 7 CPU core 176
+t-20260528052235-lvfnr-worker-0:10306:10460 [7] NCCL INFO [Proxy Service UDS] Device 7 CPU core 178
+t-20260528052235-lvfnr-worker-0:10302:10461 [3] NCCL INFO [Proxy Service] Device 3 CPU core 70
+t-20260528052235-lvfnr-worker-0:10302:10462 [3] NCCL INFO [Proxy Service UDS] Device 3 CPU core 72
+t-20260528052235-lvfnr-worker-0:10305:10463 [6] NCCL INFO [Proxy Service] Device 6 CPU core 136
+t-20260528052235-lvfnr-worker-0:10305:10464 [6] NCCL INFO [Proxy Service UDS] Device 6 CPU core 138
+t-20260528052235-lvfnr-worker-0:10304:10465 [5] NCCL INFO [Proxy Service] Device 5 CPU core 92
+t-20260528052235-lvfnr-worker-0:10304:10466 [5] NCCL INFO [Proxy Service UDS] Device 5 CPU core 96
+t-20260528052235-lvfnr-worker-0:10304:10381 [5] NCCL INFO threadThresholds 8/8/64 | 64/8/64 | 512 | 512
+t-20260528052235-lvfnr-worker-0:10304:10381 [5] NCCL INFO 24 coll channels, 24 collnet channels, 16 nvls channels, 32 p2p channels, 32 p2p channels per peer
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO threadThresholds 8/8/64 | 64/8/64 | 512 | 512
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO 24 coll channels, 24 collnet channels, 16 nvls channels, 32 p2p channels, 32 p2p channels per peer
+t-20260528052235-lvfnr-worker-0:10300:10375 [1] NCCL INFO threadThresholds 8/8/64 | 64/8/64 | 512 | 512
+t-20260528052235-lvfnr-worker-0:10300:10375 [1] NCCL INFO 24 coll channels, 24 collnet channels, 16 nvls channels, 32 p2p channels, 32 p2p channels per peer
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO CC Off, workFifoBytes 1048576
+t-20260528052235-lvfnr-worker-0:10301:10382 [2] NCCL INFO threadThresholds 8/8/64 | 64/8/64 | 512 | 512
+t-20260528052235-lvfnr-worker-0:10301:10382 [2] NCCL INFO 24 coll channels, 24 collnet channels, 16 nvls channels, 32 p2p channels, 32 p2p channels per peer
+t-20260528052235-lvfnr-worker-0:10306:10372 [7] NCCL INFO threadThresholds 8/8/64 | 64/8/64 | 512 | 512
+t-20260528052235-lvfnr-worker-0:10306:10372 [7] NCCL INFO 24 coll channels, 24 collnet channels, 16 nvls channels, 32 p2p channels, 32 p2p channels per peer
+t-20260528052235-lvfnr-worker-0:10305:10374 [6] NCCL INFO threadThresholds 8/8/64 | 64/8/64 | 512 | 512
+t-20260528052235-lvfnr-worker-0:10305:10374 [6] NCCL INFO 24 coll channels, 24 collnet channels, 16 nvls channels, 32 p2p channels, 32 p2p channels per peer
+t-20260528052235-lvfnr-worker-0:10303:10377 [4] NCCL INFO threadThresholds 8/8/64 | 64/8/64 | 512 | 512
+t-20260528052235-lvfnr-worker-0:10303:10377 [4] NCCL INFO 24 coll channels, 24 collnet channels, 16 nvls channels, 32 p2p channels, 32 p2p channels per peer
+t-20260528052235-lvfnr-worker-0:10302:10373 [3] NCCL INFO threadThresholds 8/8/64 | 64/8/64 | 512 | 512
+t-20260528052235-lvfnr-worker-0:10302:10373 [3] NCCL INFO 24 coll channels, 24 collnet channels, 16 nvls channels, 32 p2p channels, 32 p2p channels per peer
+t-20260528052235-lvfnr-worker-0:10306:10372 [7] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v4 symbol.
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v4 symbol.
+t-20260528052235-lvfnr-worker-0:10302:10373 [3] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v4 symbol.
+t-20260528052235-lvfnr-worker-0:10305:10374 [6] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v4 symbol.
+t-20260528052235-lvfnr-worker-0:10306:10372 [7] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v3 symbol.
+t-20260528052235-lvfnr-worker-0:10300:10375 [1] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v4 symbol.
+t-20260528052235-lvfnr-worker-0:10301:10382 [2] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v4 symbol.
+t-20260528052235-lvfnr-worker-0:10305:10374 [6] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v3 symbol.
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v3 symbol.
+t-20260528052235-lvfnr-worker-0:10304:10381 [5] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v4 symbol.
+t-20260528052235-lvfnr-worker-0:10305:10374 [6] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v2 symbol, using internal tuner instead.
+t-20260528052235-lvfnr-worker-0:10302:10373 [3] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v3 symbol.
+t-20260528052235-lvfnr-worker-0:10303:10377 [4] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v4 symbol.
+t-20260528052235-lvfnr-worker-0:10304:10381 [5] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v3 symbol.
+t-20260528052235-lvfnr-worker-0:10300:10375 [1] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v3 symbol.
+t-20260528052235-lvfnr-worker-0:10306:10372 [7] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v2 symbol, using internal tuner instead.
+t-20260528052235-lvfnr-worker-0:10301:10382 [2] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v3 symbol.
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v2 symbol, using internal tuner instead.
+t-20260528052235-lvfnr-worker-0:10302:10373 [3] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v2 symbol, using internal tuner instead.
+t-20260528052235-lvfnr-worker-0:10305:10374 [6] NCCL INFO ncclCommInitRankConfig comm 0x919e3e0 rank 6 nranks 8 cudaDev 6 nvmlDev 6 busId 73020 commId 0x1a57b0265de3f8e1 - Init COMPLETE
+t-20260528052235-lvfnr-worker-0:10301:10382 [2] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v2 symbol, using internal tuner instead.
+t-20260528052235-lvfnr-worker-0:10300:10375 [1] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v2 symbol, using internal tuner instead.
+t-20260528052235-lvfnr-worker-0:10303:10377 [4] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v3 symbol.
+t-20260528052235-lvfnr-worker-0:10304:10381 [5] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v2 symbol, using internal tuner instead.
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO ncclCommInitRankConfig comm 0x95c7e40 rank 0 nranks 8 cudaDev 0 nvmlDev 0 busId 65040 commId 0x1a57b0265de3f8e1 - Init COMPLETE
+t-20260528052235-lvfnr-worker-0:10306:10372 [7] NCCL INFO ncclCommInitRankConfig comm 0xa010c30 rank 7 nranks 8 cudaDev 7 nvmlDev 7 busId 75020 commId 0x1a57b0265de3f8e1 - Init COMPLETE
+t-20260528052235-lvfnr-worker-0:10302:10373 [3] NCCL INFO ncclCommInitRankConfig comm 0x8d69240 rank 3 nranks 8 cudaDev 3 nvmlDev 3 busId 6b020 commId 0x1a57b0265de3f8e1 - Init COMPLETE
+t-20260528052235-lvfnr-worker-0:10303:10377 [4] NCCL INFO TUNER/Plugin: Failed to find ncclTunerPlugin_v2 symbol, using internal tuner instead.
+t-20260528052235-lvfnr-worker-0:10301:10382 [2] NCCL INFO ncclCommInitRankConfig comm 0xcec7380 rank 2 nranks 8 cudaDev 2 nvmlDev 2 busId 69020 commId 0x1a57b0265de3f8e1 - Init COMPLETE
+t-20260528052235-lvfnr-worker-0:10300:10375 [1] NCCL INFO ncclCommInitRankConfig comm 0xcebc1b0 rank 1 nranks 8 cudaDev 1 nvmlDev 1 busId 67020 commId 0x1a57b0265de3f8e1 - Init COMPLETE
+t-20260528052235-lvfnr-worker-0:10304:10381 [5] NCCL INFO ncclCommInitRankConfig comm 0xce66c40 rank 5 nranks 8 cudaDev 5 nvmlDev 5 busId 71020 commId 0x1a57b0265de3f8e1 - Init COMPLETE
+t-20260528052235-lvfnr-worker-0:10305:10374 [6] NCCL INFO Init timings - ncclCommInitRankConfig: rank 6 nranks 8 total 2.24 (kernels 0.24, alloc 0.45, bootstrap 0.60, allgathers 0.00, topo 0.54, graphs 0.01, connections 0.36, rest 0.03)
+t-20260528052235-lvfnr-worker-0:10299:10371 [0] NCCL INFO Init timings - ncclCommInitRankConfig: rank 0 nranks 8 total 2.33 (kernels 0.19, alloc 0.29, bootstrap 0.90, allgathers 0.00, topo 0.54, graphs 0.01, connections 0.38, rest 0.02)
+t-20260528052235-lvfnr-worker-0:10303:10377 [4] NCCL INFO ncclCommInitRankConfig comm 0xc7a6880 rank 4 nranks 8 cudaDev 4 nvmlDev 4 busId 6f020 commId 0x1a57b0265de3f8e1 - Init COMPLETE
+t-20260528052235-lvfnr-worker-0:10302:10373 [3] NCCL INFO Init timings - ncclCommInitRankConfig: rank 3 nranks 8 total 2.26 (kernels 0.22, alloc 0.48, bootstrap 0.62, allgathers 0.00, topo 0.54, graphs 0.01, connections 0.37, rest 0.03)
+t-20260528052235-lvfnr-worker-0:10306:10372 [7] NCCL INFO Init timings - ncclCommInitRankConfig: rank 7 nranks 8 total 2.26 (kernels 0.21, alloc 0.47, bootstrap 0.63, allgathers 0.00, topo 0.54, graphs 0.01, connections 0.37, rest 0.03)
+t-20260528052235-lvfnr-worker-0:10301:10382 [2] NCCL INFO Init timings - ncclCommInitRankConfig: rank 2 nranks 8 total 2.10 (kernels 0.51, alloc 0.64, bootstrap 0.00, allgathers 0.00, topo 0.54, graphs 0.01, connections 0.38, rest 0.02)
+t-20260528052235-lvfnr-worker-0:10300:10375 [1] NCCL INFO Init timings - ncclCommInitRankConfig: rank 1 nranks 8 total 2.13 (kernels 0.49, alloc 0.69, bootstrap 0.00, allgathers 0.00, topo 0.54, graphs 0.01, connections 0.38, rest 0.02)
+t-20260528052235-lvfnr-worker-0:10304:10381 [5] NCCL INFO Init timings - ncclCommInitRankConfig: rank 5 nranks 8 total 2.10 (kernels 0.49, alloc 0.66, bootstrap 0.00, allgathers 0.00, topo 0.54, graphs 0.01, connections 0.36, rest 0.04)
+t-20260528052235-lvfnr-worker-0:10303:10377 [4] NCCL INFO Init timings - ncclCommInitRankConfig: rank 4 nranks 8 total 2.12 (kernels 0.47, alloc 0.71, bootstrap 0.00, allgathers 0.00, topo 0.54, graphs 0.01, connections 0.38, rest 0.01)
+t-20260528052235-lvfnr-worker-0:10299:10468 [0] NCCL INFO Channel 00/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10299:10468 [0] NCCL INFO Channel 01/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10299:10468 [0] NCCL INFO Channel 02/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10299:10468 [0] NCCL INFO Channel 03/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10299:10468 [0] NCCL INFO Channel 04/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10299:10468 [0] NCCL INFO Channel 05/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10299:10468 [0] NCCL INFO Channel 06/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10299:10468 [0] NCCL INFO Channel 07/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10299:10468 [0] NCCL INFO Channel 08/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10299:10468 [0] NCCL INFO Channel 09/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10305:10470 [6] NCCL INFO Channel 00/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10302:10469 [3] NCCL INFO Channel 00/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10299:10468 [0] NCCL INFO Channel 10/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10306:10467 [7] NCCL INFO Channel 00/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10305:10470 [6] NCCL INFO Channel 01/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10302:10469 [3] NCCL INFO Channel 01/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10305:10470 [6] NCCL INFO Channel 02/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10302:10469 [3] NCCL INFO Channel 02/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10305:10470 [6] NCCL INFO Channel 03/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10302:10469 [3] NCCL INFO Channel 03/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10305:10470 [6] NCCL INFO Channel 04/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10302:10469 [3] NCCL INFO Channel 04/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10305:10470 [6] NCCL INFO Channel 05/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10302:10469 [3] NCCL INFO Channel 05/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10300:10471 [1] NCCL INFO Channel 00/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10305:10470 [6] NCCL INFO Channel 06/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10302:10469 [3] NCCL INFO Channel 06/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10305:10470 [6] NCCL INFO Channel 07/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10299:10468 [0] NCCL INFO Channel 11/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10305:10470 [6] NCCL INFO Channel 08/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10302:10469 [3] NCCL INFO Channel 07/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10300:10471 [1] NCCL INFO Channel 01/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10305:10470 [6] NCCL INFO Channel 09/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10302:10469 [3] NCCL INFO Channel 08/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10300:10471 [1] NCCL INFO Channel 02/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10305:10470 [6] NCCL INFO Channel 10/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10306:10467 [7] NCCL INFO Channel 01/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10305:10470 [6] NCCL INFO Channel 11/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10306:10467 [7] NCCL INFO Channel 02/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10305:10470 [6] NCCL INFO Channel 12/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10306:10467 [7] NCCL INFO Channel 03/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10305:10470 [6] NCCL INFO Channel 13/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10306:10467 [7] NCCL INFO Channel 04/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10305:10470 [6] NCCL INFO Channel 14/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10306:10467 [7] NCCL INFO Channel 05/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10305:10470 [6] NCCL INFO Channel 15/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10299:10468 [0] NCCL INFO Channel 12/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10302:10469 [3] NCCL INFO Channel 09/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10299:10468 [0] NCCL INFO Channel 13/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10302:10469 [3] NCCL INFO Channel 10/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10301:10472 [2] NCCL INFO Channel 00/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10299:10468 [0] NCCL INFO Channel 14/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10304:10474 [5] NCCL INFO Channel 00/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10302:10469 [3] NCCL INFO Channel 11/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10301:10472 [2] NCCL INFO Channel 01/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10299:10468 [0] NCCL INFO Channel 15/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10303:10473 [4] NCCL INFO Channel 00/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10300:10471 [1] NCCL INFO Channel 03/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10304:10474 [5] NCCL INFO Channel 01/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10302:10469 [3] NCCL INFO Channel 12/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10301:10472 [2] NCCL INFO Channel 02/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10299:10468 [0] NCCL INFO Channel 16/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10303:10473 [4] NCCL INFO Channel 01/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10300:10471 [1] NCCL INFO Channel 04/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10304:10474 [5] NCCL INFO Channel 02/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10302:10469 [3] NCCL INFO Channel 13/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10301:10472 [2] NCCL INFO Channel 03/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10299:10468 [0] NCCL INFO Channel 17/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10303:10473 [4] NCCL INFO Channel 02/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10300:10471 [1] NCCL INFO Channel 05/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10304:10474 [5] NCCL INFO Channel 03/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10302:10469 [3] NCCL INFO Channel 14/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10301:10472 [2] NCCL INFO Channel 04/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10299:10468 [0] NCCL INFO Channel 18/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10306:10467 [7] NCCL INFO Channel 06/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10303:10473 [4] NCCL INFO Channel 03/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10305:10470 [6] NCCL INFO Channel 16/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10304:10474 [5] NCCL INFO Channel 04/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10300:10471 [1] NCCL INFO Channel 06/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10302:10469 [3] NCCL INFO Channel 15/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10301:10472 [2] NCCL INFO Channel 05/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10299:10468 [0] NCCL INFO Channel 19/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10303:10473 [4] NCCL INFO Channel 04/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10306:10467 [7] NCCL INFO Channel 07/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10305:10470 [6] NCCL INFO Channel 17/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10304:10474 [5] NCCL INFO Channel 05/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10302:10469 [3] NCCL INFO Channel 16/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10300:10471 [1] NCCL INFO Channel 07/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10301:10472 [2] NCCL INFO Channel 06/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10299:10468 [0] NCCL INFO Channel 20/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10303:10473 [4] NCCL INFO Channel 05/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10306:10467 [7] NCCL INFO Channel 08/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10305:10470 [6] NCCL INFO Channel 18/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10304:10474 [5] NCCL INFO Channel 06/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10302:10469 [3] NCCL INFO Channel 17/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10300:10471 [1] NCCL INFO Channel 08/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10301:10472 [2] NCCL INFO Channel 07/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10299:10468 [0] NCCL INFO Channel 21/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10303:10473 [4] NCCL INFO Channel 06/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10305:10470 [6] NCCL INFO Channel 19/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10306:10467 [7] NCCL INFO Channel 09/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10304:10474 [5] NCCL INFO Channel 07/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10302:10469 [3] NCCL INFO Channel 18/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10300:10471 [1] NCCL INFO Channel 09/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10301:10472 [2] NCCL INFO Channel 08/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10299:10468 [0] NCCL INFO Channel 22/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10303:10473 [4] NCCL INFO Channel 07/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10305:10470 [6] NCCL INFO Channel 20/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10306:10467 [7] NCCL INFO Channel 10/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10304:10474 [5] NCCL INFO Channel 08/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10302:10469 [3] NCCL INFO Channel 19/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10300:10471 [1] NCCL INFO Channel 10/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10301:10472 [2] NCCL INFO Channel 09/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10299:10468 [0] NCCL INFO Channel 23/0 : 0[0] -> 1[1] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10303:10473 [4] NCCL INFO Channel 08/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10305:10470 [6] NCCL INFO Channel 21/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10306:10467 [7] NCCL INFO Channel 11/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10304:10474 [5] NCCL INFO Channel 09/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10302:10469 [3] NCCL INFO Channel 20/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10300:10471 [1] NCCL INFO Channel 11/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10301:10472 [2] NCCL INFO Channel 10/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10303:10473 [4] NCCL INFO Channel 09/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10305:10470 [6] NCCL INFO Channel 22/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10306:10467 [7] NCCL INFO Channel 12/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10304:10474 [5] NCCL INFO Channel 10/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10302:10469 [3] NCCL INFO Channel 21/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10300:10471 [1] NCCL INFO Channel 12/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10301:10472 [2] NCCL INFO Channel 11/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10303:10473 [4] NCCL INFO Channel 10/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10305:10470 [6] NCCL INFO Channel 23/0 : 6[6] -> 7[7] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10306:10467 [7] NCCL INFO Channel 13/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10304:10474 [5] NCCL INFO Channel 11/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10302:10469 [3] NCCL INFO Channel 22/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10300:10471 [1] NCCL INFO Channel 13/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10301:10472 [2] NCCL INFO Channel 12/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10303:10473 [4] NCCL INFO Channel 11/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10306:10467 [7] NCCL INFO Channel 14/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10304:10474 [5] NCCL INFO Channel 12/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10302:10469 [3] NCCL INFO Channel 23/0 : 3[3] -> 4[4] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10300:10471 [1] NCCL INFO Channel 14/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10301:10472 [2] NCCL INFO Channel 13/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10303:10473 [4] NCCL INFO Channel 12/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10306:10467 [7] NCCL INFO Channel 15/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10304:10474 [5] NCCL INFO Channel 13/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10300:10471 [1] NCCL INFO Channel 15/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10301:10472 [2] NCCL INFO Channel 14/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10303:10473 [4] NCCL INFO Channel 13/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10306:10467 [7] NCCL INFO Channel 16/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10304:10474 [5] NCCL INFO Channel 14/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10300:10471 [1] NCCL INFO Channel 16/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10301:10472 [2] NCCL INFO Channel 15/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10303:10473 [4] NCCL INFO Channel 14/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10306:10467 [7] NCCL INFO Channel 17/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10304:10474 [5] NCCL INFO Channel 15/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10300:10471 [1] NCCL INFO Channel 17/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10301:10472 [2] NCCL INFO Channel 16/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10303:10473 [4] NCCL INFO Channel 15/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10306:10467 [7] NCCL INFO Channel 18/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10304:10474 [5] NCCL INFO Channel 16/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10300:10471 [1] NCCL INFO Channel 18/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10301:10472 [2] NCCL INFO Channel 17/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10303:10473 [4] NCCL INFO Channel 16/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10306:10467 [7] NCCL INFO Channel 19/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10304:10474 [5] NCCL INFO Channel 17/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10300:10471 [1] NCCL INFO Channel 19/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10301:10472 [2] NCCL INFO Channel 18/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10303:10473 [4] NCCL INFO Channel 17/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10306:10467 [7] NCCL INFO Channel 20/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10304:10474 [5] NCCL INFO Channel 18/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10300:10471 [1] NCCL INFO Channel 20/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10301:10472 [2] NCCL INFO Channel 19/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10303:10473 [4] NCCL INFO Channel 18/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10306:10467 [7] NCCL INFO Channel 21/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10304:10474 [5] NCCL INFO Channel 19/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10300:10471 [1] NCCL INFO Channel 21/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10301:10472 [2] NCCL INFO Channel 20/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10303:10473 [4] NCCL INFO Channel 19/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10306:10467 [7] NCCL INFO Channel 22/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10304:10474 [5] NCCL INFO Channel 20/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10300:10471 [1] NCCL INFO Channel 22/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10301:10472 [2] NCCL INFO Channel 21/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10303:10473 [4] NCCL INFO Channel 20/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10306:10467 [7] NCCL INFO Channel 23/0 : 7[7] -> 0[0] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10304:10474 [5] NCCL INFO Channel 21/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10300:10471 [1] NCCL INFO Channel 23/0 : 1[1] -> 2[2] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10301:10472 [2] NCCL INFO Channel 22/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10303:10473 [4] NCCL INFO Channel 21/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10304:10474 [5] NCCL INFO Channel 22/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10301:10472 [2] NCCL INFO Channel 23/0 : 2[2] -> 3[3] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10303:10473 [4] NCCL INFO Channel 22/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10304:10474 [5] NCCL INFO Channel 23/0 : 5[5] -> 6[6] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10303:10473 [4] NCCL INFO Channel 23/0 : 4[4] -> 5[5] via P2P/CUMEM
+t-20260528052235-lvfnr-worker-0:10305:10470 [6] NCCL INFO Connected all rings, use ring PXN 0 GDR 1
+t-20260528052235-lvfnr-worker-0:10304:10474 [5] NCCL INFO Connected all rings, use ring PXN 0 GDR 1
+t-20260528052235-lvfnr-worker-0:10301:10472 [2] NCCL INFO Connected all rings, use ring PXN 0 GDR 1
+t-20260528052235-lvfnr-worker-0:10300:10471 [1] NCCL INFO Connected all rings, use ring PXN 0 GDR 1
+t-20260528052235-lvfnr-worker-0:10299:10468 [0] NCCL INFO Connected all rings, use ring PXN 0 GDR 1
+t-20260528052235-lvfnr-worker-0:10306:10467 [7] NCCL INFO Connected all rings, use ring PXN 0 GDR 1
+t-20260528052235-lvfnr-worker-0:10303:10473 [4] NCCL INFO Connected all rings, use ring PXN 0 GDR 1
+t-20260528052235-lvfnr-worker-0:10302:10469 [3] NCCL INFO Connected all rings, use ring PXN 0 GDR 1
+{
+  "data_mode": "cache",
+  "cache_path": "cache/owt_t5_stream_pack1023_ultraclean_probe80k_rowvalid_10k_seed20260527_appendeos1.pt",
+  "data_path": "/e2e-data/evad-tech-vla/wanghan58/data/small_benchmarks/langflow_2604_11748/openwebtext",
+  "tokenizer_path": "/e2e-data/evad-tech-vla/wanghan58/models/hf/t5-small/tokenizer.json",
+  "text_column": "text",
+  "pack_len": 1023,
+  "append_eos": 1,
+  "num_workers": 0,
+  "shuffle_buffer": 8192,
+  "reject_txt": "cache/online_rejected.txt",
+  "out_dir": "runs/owt_t5_ultraclean10k_len1024_C4096_to_32768_exp_d768_l12_h12_gbs512_8gpu_40k_lr3e4_20260527_212316",
+  "subset_size": 10000,
+  "resume": "",
+  "steps": 40000,
+  "batch_size": 32,
+  "grad_accum": 2,
+  "lr": 0.0003,
+  "log_every": 50,
+  "save_every": 1000,
+  "dim": 768,
+  "layers": 12,
+  "heads": 12,
+  "mlp_dim": 3072,
+  "time_tokens": 4,
+  "c_min": 4096.0,
+  "c_max": 32768.0,
+  "c_schedule": "exp",
+  "seed": 1234
+}
+[data] mode=cache rows=10000 length=1024 vocab=32100 seen=80003 dropped=15766 bos=1:</s> eos=1:</s>
+t-20260528052235-lvfnr-worker-0:10302:10563 [3] NCCL INFO NVLS comm 0x8d69240 headRank 3 nHeads 8 buffSize 1048576 nvlsPerRankSize 33554432 nvlsTotalSize 268435456
+t-20260528052235-lvfnr-worker-0:10300:10564 [1] NCCL INFO NVLS comm 0xcebc1b0 headRank 1 nHeads 8 buffSize 1048576 nvlsPerRankSize 33554432 nvlsTotalSize 268435456
+t-20260528052235-lvfnr-worker-0:10306:10565 [7] NCCL INFO NVLS comm 0xa010c30 headRank 7 nHeads 8 buffSize 1048576 nvlsPerRankSize 33554432 nvlsTotalSize 268435456
+t-20260528052235-lvfnr-worker-0:10303:10566 [4] NCCL INFO NVLS comm 0xc7a6880 headRank 4 nHeads 8 buffSize 1048576 nvlsPerRankSize 33554432 nvlsTotalSize 268435456
+t-20260528052235-lvfnr-worker-0:10301:10567 [2] NCCL INFO NVLS comm 0xcec7380 headRank 2 nHeads 8 buffSize 1048576 nvlsPerRankSize 33554432 nvlsTotalSize 268435456
+t-20260528052235-lvfnr-worker-0:10304:10568 [5] NCCL INFO NVLS comm 0xce66c40 headRank 5 nHeads 8 buffSize 1048576 nvlsPerRankSize 33554432 nvlsTotalSize 268435456
+t-20260528052235-lvfnr-worker-0:10305:10570 [6] NCCL INFO NVLS comm 0x919e3e0 headRank 6 nHeads 8 buffSize 1048576 nvlsPerRankSize 33554432 nvlsTotalSize 268435456
+t-20260528052235-lvfnr-worker-0:10299:10569 [0] NCCL INFO NVLS comm 0x95c7e40 headRank 0 nHeads 8 buffSize 1048576 nvlsPerRankSize 33554432 nvlsTotalSize 268435456
+step=50 loss=7.1228 {'pos0_bos_p': 0.0019045539665967226, 'pos0_bos_top1': 0, 'last_eos_p': 0.010252108797430992, 'last_eos_top1': 0}
+step=100 loss=7.1434 {'pos0_bos_p': 0.00042073283111676574, 'pos0_bos_top1': 0, 'last_eos_p': 0.9972795844078064, 'last_eos_top1': 4}
+step=150 loss=6.7899 {'pos0_bos_p': 0.0011103940196335316, 'pos0_bos_top1': 0, 'last_eos_p': 0.9403679966926575, 'last_eos_top1': 4}
+step=200 loss=6.5186 {'pos0_bos_p': 0.0008917325758375227, 'pos0_bos_top1': 0, 'last_eos_p': 0.9825672507286072, 'last_eos_top1': 4}
+step=250 loss=6.0590 {'pos0_bos_p': 0.001217335811816156, 'pos0_bos_top1': 0, 'last_eos_p': 0.9862273335456848, 'last_eos_top1': 4}
+step=300 loss=6.1027 {'pos0_bos_p': 0.002536393003538251, 'pos0_bos_top1': 0, 'last_eos_p': 0.9804686307907104, 'last_eos_top1': 4}
+step=350 loss=5.8705 {'pos0_bos_p': 0.001197237172164023, 'pos0_bos_top1': 0, 'last_eos_p': 0.9804034233093262, 'last_eos_top1': 4}
+step=400 loss=5.7510 {'pos0_bos_p': 0.001667062402702868, 'pos0_bos_top1': 0, 'last_eos_p': 0.9864185452461243, 'last_eos_top1': 4}
+step=450 loss=5.1534 {'pos0_bos_p': 0.0024044252932071686, 'pos0_bos_top1': 0, 'last_eos_p': 0.9860549569129944, 'last_eos_top1': 4}
+step=500 loss=5.1372 {'pos0_bos_p': 0.001829120796173811, 'pos0_bos_top1': 0, 'last_eos_p': 0.9869738817214966, 'last_eos_top1': 4}
+step=550 loss=5.2508 {'pos0_bos_p': 0.0011939660180360079, 'pos0_bos_top1': 0, 'last_eos_p': 0.9870648384094238, 'last_eos_top1': 4}
+step=600 loss=4.7833 {'pos0_bos_p': 0.001596515066921711, 'pos0_bos_top1': 0, 'last_eos_p': 0.9869128465652466, 'last_eos_top1': 4}
+step=650 loss=4.8124 {'pos0_bos_p': 0.001284916652366519, 'pos0_bos_top1': 0, 'last_eos_p': 0.9870219826698303, 'last_eos_top1': 4}
+step=700 loss=4.6599 {'pos0_bos_p': 0.0014337411848828197, 'pos0_bos_top1': 0, 'last_eos_p': 0.9863752722740173, 'last_eos_top1': 4}
+step=750 loss=4.4790 {'pos0_bos_p': 0.0012800463009625673, 'pos0_bos_top1': 0, 'last_eos_p': 0.987720787525177, 'last_eos_top1': 4}
+step=800 loss=4.5593 {'pos0_bos_p': 0.0014664315385743976, 'pos0_bos_top1': 0, 'last_eos_p': 0.9885030388832092, 'last_eos_top1': 4}
+step=850 loss=4.4832 {'pos0_bos_p': 0.0012145357904955745, 'pos0_bos_top1': 0, 'last_eos_p': 0.9880351424217224, 'last_eos_top1': 4}
+step=900 loss=4.7124 {'pos0_bos_p': 0.0010197763331234455, 'pos0_bos_top1': 0, 'last_eos_p': 0.9898796081542969, 'last_eos_top1': 4}
+step=950 loss=4.6711 {'pos0_bos_p': 0.0011980183189734817, 'pos0_bos_top1': 0, 'last_eos_p': 0.9903967380523682, 'last_eos_top1': 4}
+step=1000 loss=4.3427 {'pos0_bos_p': 0.0012490765657275915, 'pos0_bos_top1': 0, 'last_eos_p': 0.9919277429580688, 'last_eos_top1': 4}
+step=1050 loss=4.4249 {'pos0_bos_p': 0.0015577699523419142, 'pos0_bos_top1': 0, 'last_eos_p': 0.9943477511405945, 'last_eos_top1': 4}
+step=1100 loss=4.4758 {'pos0_bos_p': 0.0007777540595270693, 'pos0_bos_top1': 0, 'last_eos_p': 0.9934155941009521, 'last_eos_top1': 4}
+step=1150 loss=3.7029 {'pos0_bos_p': 0.0009281560778617859, 'pos0_bos_top1': 0, 'last_eos_p': 0.9941577911376953, 'last_eos_top1': 4}
+step=1200 loss=4.6245 {'pos0_bos_p': 0.0009716648492030799, 'pos0_bos_top1': 0, 'last_eos_p': 0.9952744245529175, 'last_eos_top1': 4}
+step=1250 loss=3.6213 {'pos0_bos_p': 0.0007597565418109298, 'pos0_bos_top1': 0, 'last_eos_p': 0.995770275592804, 'last_eos_top1': 4}
+step=1300 loss=3.6857 {'pos0_bos_p': 0.0007579440716654062, 'pos0_bos_top1': 0, 'last_eos_p': 0.9962783455848694, 'last_eos_top1': 4}
+step=1350 loss=3.8707 {'pos0_bos_p': 0.0007095331093296409, 'pos0_bos_top1': 0, 'last_eos_p': 0.99689120054245, 'last_eos_top1': 4}
+step=1400 loss=3.7336 {'pos0_bos_p': 0.0008909617899917066, 'pos0_bos_top1': 0, 'last_eos_p': 0.9970274567604065, 'last_eos_top1': 4}
+step=1450 loss=3.6988 {'pos0_bos_p': 0.0009782277047634125, 'pos0_bos_top1': 0, 'last_eos_p': 0.9973872303962708, 'last_eos_top1': 4}
+step=1500 loss=3.5358 {'pos0_bos_p': 0.0009436120162717998, 'pos0_bos_top1': 0, 'last_eos_p': 0.9975798726081848, 'last_eos_top1': 4}
+step=1550 loss=4.3268 {'pos0_bos_p': 0.0009957518195733428, 'pos0_bos_top1': 0, 'last_eos_p': 0.9976412057876587, 'last_eos_top1': 4}
+step=1600 loss=3.4677 {'pos0_bos_p': 0.000908614369109273, 'pos0_bos_top1': 0, 'last_eos_p': 0.9979562759399414, 'last_eos_top1': 4}
+step=1650 loss=3.5105 {'pos0_bos_p': 0.0009317692019976676, 'pos0_bos_top1': 0, 'last_eos_p': 0.998079776763916, 'last_eos_top1': 4}
+step=1700 loss=3.7310 {'pos0_bos_p': 0.0010493824956938624, 'pos0_bos_top1': 0, 'last_eos_p': 0.9982014894485474, 'last_eos_top1': 4}
+step=1750 loss=3.6179 {'pos0_bos_p': 0.0009264908730983734, 'pos0_bos_top1': 0, 'last_eos_p': 0.9982447624206543, 'last_eos_top1': 4}
+step=1800 loss=3.7720 {'pos0_bos_p': 0.0010388639057055116, 'pos0_bos_top1': 0, 'last_eos_p': 0.9982740879058838, 'last_eos_top1': 4}
+step=1850 loss=2.9854 {'pos0_bos_p': 0.0010407152585685253, 'pos0_bos_top1': 0, 'last_eos_p': 0.998422384262085, 'last_eos_top1': 4}
+step=1900 loss=3.4438 {'pos0_bos_p': 0.0011615091934800148, 'pos0_bos_top1': 0, 'last_eos_p': 0.9984297156333923, 'last_eos_top1': 4}
+step=1950 loss=3.0192 {'pos0_bos_p': 0.0009644241072237492, 'pos0_bos_top1': 0, 'last_eos_p': 0.9983976483345032, 'last_eos_top1': 4}
+step=2000 loss=2.9811 {'pos0_bos_p': 0.0010724756866693497, 'pos0_bos_top1': 0, 'last_eos_p': 0.9983878135681152, 'last_eos_top1': 4}
+step=2050 loss=2.9872 {'pos0_bos_p': 0.0007766140624880791, 'pos0_bos_top1': 0, 'last_eos_p': 0.9984954595565796, 'last_eos_top1': 4}
+step=2100 loss=3.5225 {'pos0_bos_p': 0.000930797541514039, 'pos0_bos_top1': 0, 'last_eos_p': 0.9985345602035522, 'last_eos_top1': 4}
+step=2150 loss=3.1885 {'pos0_bos_p': 0.000730251194909215, 'pos0_bos_top1': 0, 'last_eos_p': 0.9986714124679565, 'last_eos_top1': 4}
+step=2200 loss=3.4761 {'pos0_bos_p': 0.0009825056185945868, 'pos0_bos_top1': 0, 'last_eos_p': 0.9985960125923157, 'last_eos_top1': 4}
+step=2250 loss=2.9302 {'pos0_bos_p': 0.001204384258016944, 'pos0_bos_top1': 0, 'last_eos_p': 0.9987609386444092, 'last_eos_top1': 4}
+step=2300 loss=3.3332 {'pos0_bos_p': 0.0009695333428680897, 'pos0_bos_top1': 0, 'last_eos_p': 0.9987534284591675, 'last_eos_top1': 4}
+step=2350 loss=3.0643 {'pos0_bos_p': 0.0009657082264311612, 'pos0_bos_top1': 0, 'last_eos_p': 0.998688280582428, 'last_eos_top1': 4}
+step=2400 loss=3.0399 {'pos0_bos_p': 0.0012310727033764124, 'pos0_bos_top1': 0, 'last_eos_p': 0.9988003969192505, 'last_eos_top1': 4}
+step=2450 loss=3.4925 {'pos0_bos_p': 0.0008470729226246476, 'pos0_bos_top1': 0, 'last_eos_p': 0.998854398727417, 'last_eos_top1': 4}
+step=2500 loss=3.2871 {'pos0_bos_p': 0.0006296156789176166, 'pos0_bos_top1': 0, 'last_eos_p': 0.9987519979476929, 'last_eos_top1': 4}
+step=2550 loss=3.1758 {'pos0_bos_p': 0.0009265568223781884, 'pos0_bos_top1': 0, 'last_eos_p': 0.9988577365875244, 'last_eos_top1': 4}
+step=2600 loss=2.8663 {'pos0_bos_p': 0.0008602665038779378, 'pos0_bos_top1': 0, 'last_eos_p': 0.9987666606903076, 'last_eos_top1': 4}
+step=2650 loss=3.0329 {'pos0_bos_p': 0.0008244539494626224, 'pos0_bos_top1': 0, 'last_eos_p': 0.9988536834716797, 'last_eos_top1': 4}
+step=2700 loss=3.4154 {'pos0_bos_p': 0.0007191206677816808, 'pos0_bos_top1': 0, 'last_eos_p': 0.998994767665863, 'last_eos_top1': 4}
+step=2750 loss=3.3015 {'pos0_bos_p': 0.0007164740818552673, 'pos0_bos_top1': 0, 'last_eos_p': 0.9991008043289185, 'last_eos_top1': 4}
+step=2800 loss=2.6920 {'pos0_bos_p': 0.0008439717930741608, 'pos0_bos_top1': 0, 'last_eos_p': 0.9990654587745667, 'last_eos_top1': 4}
+step=2850 loss=3.0997 {'pos0_bos_p': 0.0008432737668044865, 'pos0_bos_top1': 0, 'last_eos_p': 0.9990019202232361, 'last_eos_top1': 4}
+step=2900 loss=3.0128 {'pos0_bos_p': 0.0008241361356340349, 'pos0_bos_top1': 0, 'last_eos_p': 0.9991840720176697, 'last_eos_top1': 4}
+step=2950 loss=2.3291 {'pos0_bos_p': 0.0007474969606846571, 'pos0_bos_top1': 0, 'last_eos_p': 0.9992383718490601, 'last_eos_top1': 4}
+step=3000 loss=3.3041 {'pos0_bos_p': 0.0008691821130923927, 'pos0_bos_top1': 0, 'last_eos_p': 0.9991764426231384, 'last_eos_top1': 4}
+step=3050 loss=2.6755 {'pos0_bos_p': 0.000884750799741596, 'pos0_bos_top1': 0, 'last_eos_p': 0.9991933703422546, 'last_eos_top1': 4}
+step=3100 loss=2.9095 {'pos0_bos_p': 0.0007487277616746724, 'pos0_bos_top1': 0, 'last_eos_p': 0.9991087317466736, 'last_eos_top1': 4}
+step=3150 loss=2.6376 {'pos0_bos_p': 0.0008386721601709723, 'pos0_bos_top1': 0, 'last_eos_p': 0.999219536781311, 'last_eos_top1': 4}
+step=3200 loss=3.0389 {'pos0_bos_p': 0.0005033225752413273, 'pos0_bos_top1': 0, 'last_eos_p': 0.9991424083709717, 'last_eos_top1': 4}
+step=3250 loss=3.5516 {'pos0_bos_p': 0.0007498532067984343, 'pos0_bos_top1': 0, 'last_eos_p': 0.9991549253463745, 'last_eos_top1': 4}
+step=3300 loss=2.9669 {'pos0_bos_p': 0.0006075959536246955, 'pos0_bos_top1': 0, 'last_eos_p': 0.9993537068367004, 'last_eos_top1': 4}
+step=3350 loss=2.8941 {'pos0_bos_p': 0.0006255818880163133, 'pos0_bos_top1': 0, 'last_eos_p': 0.9992604851722717, 'last_eos_top1': 4}
+step=3400 loss=2.4231 {'pos0_bos_p': 0.0006005432223901153, 'pos0_bos_top1': 0, 'last_eos_p': 0.9993374943733215, 'last_eos_top1': 4}
+step=3450 loss=3.0347 {'pos0_bos_p': 0.000714322435669601, 'pos0_bos_top1': 0, 'last_eos_p': 0.9993426203727722, 'last_eos_top1': 4}
+step=3500 loss=2.6405 {'pos0_bos_p': 0.0006911610253155231, 'pos0_bos_top1': 0, 'last_eos_p': 0.999247670173645, 'last_eos_top1': 4}
+step=3550 loss=3.0160 {'pos0_bos_p': 0.0006793943466618657, 'pos0_bos_top1': 0, 'last_eos_p': 0.9993601441383362, 'last_eos_top1': 4}
+step=3600 loss=2.7247 {'pos0_bos_p': 0.0006277165375649929, 'pos0_bos_top1': 0, 'last_eos_p': 0.999326229095459, 'last_eos_top1': 4}
+step=3650 loss=2.6101 {'pos0_bos_p': 0.0007806556532159448, 'pos0_bos_top1': 0, 'last_eos_p': 0.9992813467979431, 'last_eos_top1': 4}
+step=3700 loss=2.7879 {'pos0_bos_p': 0.000695927650667727, 'pos0_bos_top1': 0, 'last_eos_p': 0.9992989301681519, 'last_eos_top1': 4}
+step=3750 loss=3.1240 {'pos0_bos_p': 0.0004254880186636001, 'pos0_bos_top1': 0, 'last_eos_p': 0.9991540908813477, 'last_eos_top1': 4}
+step=3800 loss=3.0976 {'pos0_bos_p': 0.0008554062806069851, 'pos0_bos_top1': 0, 'last_eos_p': 0.9992836117744446, 'last_eos_top1': 4}
+step=3850 loss=3.0993 {'pos0_bos_p': 0.0004907391266897321, 'pos0_bos_top1': 0, 'last_eos_p': 0.9993964433670044, 'last_eos_top1': 4}
+step=3900 loss=2.9258 {'pos0_bos_p': 0.0006513699190691113, 'pos0_bos_top1': 0, 'last_eos_p': 0.9992951154708862, 'last_eos_top1': 4}
+step=3950 loss=3.1170 {'pos0_bos_p': 0.0007697808905504644, 'pos0_bos_top1': 0, 'last_eos_p': 0.9992783665657043, 'last_eos_top1': 4}
+step=4000 loss=2.8308 {'pos0_bos_p': 0.0004717964620795101, 'pos0_bos_top1': 0, 'last_eos_p': 0.999384880065918, 'last_eos_top1': 4}
+step=4050 loss=2.8822 {'pos0_bos_p': 0.000635715841781348, 'pos0_bos_top1': 0, 'last_eos_p': 0.9994264841079712, 'last_eos_top1': 4}
+step=4100 loss=3.0471 {'pos0_bos_p': 0.000768937636166811, 'pos0_bos_top1': 0, 'last_eos_p': 0.9993540644645691, 'last_eos_top1': 4}
+step=4150 loss=2.7838 {'pos0_bos_p': 0.0008931310730986297, 'pos0_bos_top1': 0, 'last_eos_p': 0.9992462396621704, 'last_eos_top1': 4}
+step=4200 loss=2.7448 {'pos0_bos_p': 0.0005247585359029472, 'pos0_bos_top1': 0, 'last_eos_p': 0.9993844032287598, 'last_eos_top1': 4}
+step=4250 loss=3.0325 {'pos0_bos_p': 0.0006107465014792979, 'pos0_bos_top1': 0, 'last_eos_p': 0.9994292855262756, 'last_eos_top1': 4}
+step=4300 loss=2.9740 {'pos0_bos_p': 0.0006588331889361143, 'pos0_bos_top1': 0, 'last_eos_p': 0.999228835105896, 'last_eos_top1': 4}
+step=4350 loss=3.0795 {'pos0_bos_p': 0.0007223390857689083, 'pos0_bos_top1': 0, 'last_eos_p': 0.999248206615448, 'last_eos_top1': 4}
+step=4400 loss=2.7557 {'pos0_bos_p': 0.0009390123886987567, 'pos0_bos_top1': 0, 'last_eos_p': 0.9993976354598999, 'last_eos_top1': 4}
+step=4450 loss=2.7971 {'pos0_bos_p': 0.0007036151364445686, 'pos0_bos_top1': 0, 'last_eos_p': 0.9995077848434448, 'last_eos_top1': 4}
+step=4500 loss=2.5748 {'pos0_bos_p': 0.000773977255448699, 'pos0_bos_top1': 0, 'last_eos_p': 0.9993659853935242, 'last_eos_top1': 4}
+step=4550 loss=2.9357 {'pos0_bos_p': 0.0008800712530501187, 'pos0_bos_top1': 0, 'last_eos_p': 0.9993867874145508, 'last_eos_top1': 4}
+step=4600 loss=2.5206 {'pos0_bos_p': 0.0007377828587777913, 'pos0_bos_top1': 0, 'last_eos_p': 0.9994648098945618, 'last_eos_top1': 4}
+step=4650 loss=2.7212 {'pos0_bos_p': 0.0007580575183965266, 'pos0_bos_top1': 0, 'last_eos_p': 0.9994601607322693, 'last_eos_top1': 4}
+step=4700 loss=2.8740 {'pos0_bos_p': 0.0006156228482723236, 'pos0_bos_top1': 0, 'last_eos_p': 0.999264657497406, 'last_eos_top1': 4}
+step=4750 loss=3.0526 {'pos0_bos_p': 0.0006342521519400179, 'pos0_bos_top1': 0, 'last_eos_p': 0.999297022819519, 'last_eos_top1': 4}
+step=4800 loss=2.8730 {'pos0_bos_p': 0.000740743416827172, 'pos0_bos_top1': 0, 'last_eos_p': 0.9992921352386475, 'last_eos_top1': 4}
+step=4850 loss=2.8685 {'pos0_bos_p': 0.0006380485720001161, 'pos0_bos_top1': 0, 'last_eos_p': 0.9992501139640808, 'last_eos_top1': 4}
+step=4900 loss=2.7638 {'pos0_bos_p': 0.0005786812398582697, 'pos0_bos_top1': 0, 'last_eos_p': 0.9991870522499084, 'last_eos_top1': 4}
+step=4950 loss=2.6238 {'pos0_bos_p': 0.0006091449176892638, 'pos0_bos_top1': 0, 'last_eos_p': 0.9991648197174072, 'last_eos_top1': 4}
+step=5000 loss=2.4685 {'pos0_bos_p': 0.000684699509292841, 'pos0_bos_top1': 0, 'last_eos_p': 0.9992815852165222, 'last_eos_top1': 4}
+step=5050 loss=2.4113 {'pos0_bos_p': 0.00050110905431211, 'pos0_bos_top1': 0, 'last_eos_p': 0.999278724193573, 'last_eos_top1': 4}
+step=5100 loss=2.9088 {'pos0_bos_p': 0.0005313889123499393, 'pos0_bos_top1': 0, 'last_eos_p': 0.999170184135437, 'last_eos_top1': 4}
+step=5150 loss=2.5220 {'pos0_bos_p': 0.0005848270375281572, 'pos0_bos_top1': 0, 'last_eos_p': 0.9991790652275085, 'last_eos_top1': 4}
+step=5200 loss=2.5151 {'pos0_bos_p': 0.0008432670729234815, 'pos0_bos_top1': 0, 'last_eos_p': 0.9993453621864319, 'last_eos_top1': 4}
+step=5250 loss=2.8306 {'pos0_bos_p': 0.0009108633967116475, 'pos0_bos_top1': 0, 'last_eos_p': 0.9992652535438538, 'last_eos_top1': 4}
+step=5300 loss=2.5167 {'pos0_bos_p': 0.0010536877671256661, 'pos0_bos_top1': 0, 'last_eos_p': 0.9991229176521301, 'last_eos_top1': 4}
+step=5350 loss=2.8203 {'pos0_bos_p': 0.0006855487590655684, 'pos0_bos_top1': 0, 'last_eos_p': 0.9993834495544434, 'last_eos_top1': 4}
+step=5400 loss=2.5030 {'pos0_bos_p': 0.0006188146653585136, 'pos0_bos_top1': 0, 'last_eos_p': 0.9992848038673401, 'last_eos_top1': 4}
+step=5450 loss=2.6654 {'pos0_bos_p': 0.0006890347576700151, 'pos0_bos_top1': 0, 'last_eos_p': 0.999160647392273, 'last_eos_top1': 4}
+step=5500 loss=2.3946 {'pos0_bos_p': 0.000737423135433346, 'pos0_bos_top1': 0, 'last_eos_p': 0.9991815686225891, 'last_eos_top1': 4}
+step=5550 loss=3.0131 {'pos0_bos_p': 0.0008019988308660686, 'pos0_bos_top1': 0, 'last_eos_p': 0.9990608096122742, 'last_eos_top1': 4}
+step=5600 loss=2.7925 {'pos0_bos_p': 0.0005938691319897771, 'pos0_bos_top1': 0, 'last_eos_p': 0.9991318583488464, 'last_eos_top1': 4}
+step=5650 loss=2.7450 {'pos0_bos_p': 0.0006745798164047301, 'pos0_bos_top1': 0, 'last_eos_p': 0.9989751577377319, 'last_eos_top1': 4}
+step=5700 loss=2.6491 {'pos0_bos_p': 0.0008223381591960788, 'pos0_bos_top1': 0, 'last_eos_p': 0.999255359172821, 'last_eos_top1': 4}
+step=5750 loss=2.8562 {'pos0_bos_p': 0.0006917612627148628, 'pos0_bos_top1': 0, 'last_eos_p': 0.999447762966156, 'last_eos_top1': 4}
+step=5800 loss=2.5180 {'pos0_bos_p': 0.0008420003578066826, 'pos0_bos_top1': 0, 'last_eos_p': 0.9990509152412415, 'last_eos_top1': 4}
+step=5850 loss=2.8989 {'pos0_bos_p': 0.0008550017373636365, 'pos0_bos_top1': 0, 'last_eos_p': 0.9990940093994141, 'last_eos_top1': 4}
+step=5900 loss=1.9623 {'pos0_bos_p': 0.0008237525471486151, 'pos0_bos_top1': 0, 'last_eos_p': 0.9991647005081177, 'last_eos_top1': 4}
+step=5950 loss=3.0246 {'pos0_bos_p': 0.0010133307659998536, 'pos0_bos_top1': 0, 'last_eos_p': 0.999256432056427, 'last_eos_top1': 4}
+step=6000 loss=2.4829 {'pos0_bos_p': 0.0009167233947664499, 'pos0_bos_top1': 0, 'last_eos_p': 0.999213695526123, 'last_eos_top1': 4}
+step=6050 loss=2.7540 {'pos0_bos_p': 0.000716303416993469, 'pos0_bos_top1': 0, 'last_eos_p': 0.9991413354873657, 'last_eos_top1': 4}
+step=6100 loss=3.2656 {'pos0_bos_p': 0.000814702536445111, 'pos0_bos_top1': 0, 'last_eos_p': 0.9991452693939209, 'last_eos_top1': 4}
+step=6150 loss=2.6433 {'pos0_bos_p': 0.0007208568858914077, 'pos0_bos_top1': 0, 'last_eos_p': 0.9992269277572632, 'last_eos_top1': 4}
+step=6200 loss=2.3956 {'pos0_bos_p': 0.0006785871810279787, 'pos0_bos_top1': 0, 'last_eos_p': 0.9992161989212036, 'last_eos_top1': 4}
+step=6250 loss=2.4800 {'pos0_bos_p': 0.000884954642970115, 'pos0_bos_top1': 0, 'last_eos_p': 0.9991790652275085, 'last_eos_top1': 4}
+step=6300 loss=2.8687 {'pos0_bos_p': 0.0006654368480667472, 'pos0_bos_top1': 0, 'last_eos_p': 0.9989078044891357, 'last_eos_top1': 4}
+step=6350 loss=2.7883 {'pos0_bos_p': 0.0007044913363642991, 'pos0_bos_top1': 0, 'last_eos_p': 0.9985566735267639, 'last_eos_top1': 4}
+step=6400 loss=2.6162 {'pos0_bos_p': 0.000603503140155226, 'pos0_bos_top1': 0, 'last_eos_p': 0.9987297654151917, 'last_eos_top1': 4}
+step=6450 loss=2.6558 {'pos0_bos_p': 0.0006924325134605169, 'pos0_bos_top1': 0, 'last_eos_p': 0.9990050196647644, 'last_eos_top1': 4}
+step=6500 loss=2.8178 {'pos0_bos_p': 0.0008079643012024462, 'pos0_bos_top1': 0, 'last_eos_p': 0.9986687898635864, 'last_eos_top1': 4}
+step=6550 loss=2.3740 {'pos0_bos_p': 0.0008527531754225492, 'pos0_bos_top1': 0, 'last_eos_p': 0.9989016056060791, 'last_eos_top1': 4}
+step=6600 loss=2.5708 {'pos0_bos_p': 0.0009096905705519021, 'pos0_bos_top1': 0, 'last_eos_p': 0.9988553524017334, 'last_eos_top1': 4}
+step=6650 loss=2.3440 {'pos0_bos_p': 0.0008201909367926419, 'pos0_bos_top1': 0, 'last_eos_p': 0.9989750385284424, 'last_eos_top1': 4}
+step=6700 loss=2.4336 {'pos0_bos_p': 0.0008251758990809321, 'pos0_bos_top1': 0, 'last_eos_p': 0.9988362193107605, 'last_eos_top1': 4}
+step=6750 loss=2.7569 {'pos0_bos_p': 0.0007772842654958367, 'pos0_bos_top1': 0, 'last_eos_p': 0.9991196990013123, 'last_eos_top1': 4}
+step=6800 loss=2.2757 {'pos0_bos_p': 0.0005473861237987876, 'pos0_bos_top1': 0, 'last_eos_p': 0.9990571141242981, 'last_eos_top1': 4}
+step=6850 loss=2.6205 {'pos0_bos_p': 0.0008291201665997505, 'pos0_bos_top1': 0, 'last_eos_p': 0.9987955093383789, 'last_eos_top1': 4}
+step=6900 loss=2.1564 {'pos0_bos_p': 0.0008855836349539459, 'pos0_bos_top1': 0, 'last_eos_p': 0.9991982579231262, 'last_eos_top1': 4}
+step=6950 loss=2.6010 {'pos0_bos_p': 0.0006183949881233275, 'pos0_bos_top1': 0, 'last_eos_p': 0.9990167617797852, 'last_eos_top1': 4}
+step=7000 loss=2.8135 {'pos0_bos_p': 0.000769881356973201, 'pos0_bos_top1': 0, 'last_eos_p': 0.9988486766815186, 'last_eos_top1': 4}
+step=7050 loss=2.5095 {'pos0_bos_p': 0.0006667476845905185, 'pos0_bos_top1': 0, 'last_eos_p': 0.9989122152328491, 'last_eos_top1': 4}
+step=7100 loss=2.4432 {'pos0_bos_p': 0.0008150693029165268, 'pos0_bos_top1': 0, 'last_eos_p': 0.9988364577293396, 'last_eos_top1': 4}
+step=7150 loss=2.7365 {'pos0_bos_p': 0.000947561115026474, 'pos0_bos_top1': 0, 'last_eos_p': 0.9986352324485779, 'last_eos_top1': 4}
+step=7200 loss=2.0627 {'pos0_bos_p': 0.0007095952169038355, 'pos0_bos_top1': 0, 'last_eos_p': 0.9989629983901978, 'last_eos_top1': 4}
+step=7250 loss=2.6835 {'pos0_bos_p': 0.0006409320048987865, 'pos0_bos_top1': 0, 'last_eos_p': 0.998291552066803, 'last_eos_top1': 4}
+step=7300 loss=2.3597 {'pos0_bos_p': 0.0008239856106229126, 'pos0_bos_top1': 0, 'last_eos_p': 0.9990217685699463, 'last_eos_top1': 4}
+step=7350 loss=2.8695 {'pos0_bos_p': 0.0007445314549840987, 'pos0_bos_top1': 0, 'last_eos_p': 0.9984349608421326, 'last_eos_top1': 4}
+step=7400 loss=2.6584 {'pos0_bos_p': 0.0008714564028196037, 'pos0_bos_top1': 0, 'last_eos_p': 0.9987302422523499, 'last_eos_top1': 4}
+step=7450 loss=2.7820 {'pos0_bos_p': 0.0007681816932745278, 'pos0_bos_top1': 0, 'last_eos_p': 0.998702883720398, 'last_eos_top1': 4}
+step=7500 loss=3.0568 {'pos0_bos_p': 0.0007770757074467838, 'pos0_bos_top1': 0, 'last_eos_p': 0.998542070388794, 'last_eos_top1': 4}
+step=7550 loss=2.3619 {'pos0_bos_p': 0.0006897128187119961, 'pos0_bos_top1': 0, 'last_eos_p': 0.9985713958740234, 'last_eos_top1': 4}
+step=7600 loss=2.0023 {'pos0_bos_p': 0.0008945409208536148, 'pos0_bos_top1': 0, 'last_eos_p': 0.998450517654419, 'last_eos_top1': 4}
+step=7650 loss=3.4515 {'pos0_bos_p': 0.0008103225845843554, 'pos0_bos_top1': 0, 'last_eos_p': 0.9984140396118164, 'last_eos_top1': 4}
+step=7700 loss=2.3593 {'pos0_bos_p': 0.0007089058053679764, 'pos0_bos_top1': 0, 'last_eos_p': 0.9984318614006042, 'last_eos_top1': 4}
+step=7750 loss=2.8948 {'pos0_bos_p': 0.0008973771473392844, 'pos0_bos_top1': 0, 'last_eos_p': 0.99854975938797, 'last_eos_top1': 4}
+step=7800 loss=2.2413 {'pos0_bos_p': 0.000924734165892005, 'pos0_bos_top1': 0, 'last_eos_p': 0.9985358715057373, 'last_eos_top1': 4}

LTA_openwebtext_dualt/mini_owt_logdirichlet/runs/owt_t5_elftokenized_full_len1024_C1_to_1024_pow1_d768_l12_h12_gbs512_8gpu_50ep_lr3e4_elfopt_t5embed_unfixed_stateprobadd_selfcond_ce_fast_20260531_230026/step_082000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:effbfbb0a79df5830618842d59449f38ac90c9439be2043870d67f04c04fe987
+size 515519058

LTA_openwebtext_dualt/mini_owt_logdirichlet/runs/owt_t5_elftokenized_full_len1024_C1_to_1024_pow1_d768_l12_h12_gbs512_8gpu_50ep_lr3e4_elfopt_t5embed_unfixed_stateprobadd_selfcond_ce_fast_20260531_230026/step_155000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1cdcc5b38fdcbe929e9dadf44e17a0be9f2958eb6c15ffb338cc1389fcf14e69
+size 515519058