megatron-core 0.14.0rc2__tar.gz → 0.14.0rc3__tar.gz

{megatron_core-0.14.0rc2/megatron_core.egg-info → megatron_core-0.14.0rc3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: megatron-core
-Version: 0.14.0rc2
+Version: 0.14.0rc3
 Summary: Megatron Core - a library for efficient and scalable training of transformer based models
 Author-email: NVIDIA <nemo-toolkit@nvidia.com>
 Maintainer-email: NVIDIA <nemo-toolkit@nvidia.com>
@@ -31,6 +31,7 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: torch
 Requires-Dist: numpy<2.0.0
+Requires-Dist: packaging~=25.0
 Provides-Extra: mlm
 Requires-Dist: flask-restful; extra == "mlm"
 Requires-Dist: sentencepiece; extra == "mlm"
@@ -38,14 +39,16 @@ Requires-Dist: tiktoken; extra == "mlm"
 Requires-Dist: wandb; extra == "mlm"
 Provides-Extra: dev
 Requires-Dist: tqdm; extra == "dev"
-Requires-Dist: einops; extra == "dev"
-Requires-Dist: tensorstore!=0.1.46,!=0.1.72; extra == "dev"
-Requires-Dist: nvtx; extra == "dev"
-Requires-Dist: transformers; extra == "dev"
-Requires-Dist: multi-storage-client; extra == "dev"
+Requires-Dist: einops~=0.8; extra == "dev"
+Requires-Dist: tensorstore!=0.1.46,!=0.1.72,~=0.1; extra == "dev"
+Requires-Dist: nvtx~=0.2; extra == "dev"
+Requires-Dist: transformers~=4.53; extra == "dev"
+Requires-Dist: multi-storage-client~=0.20.3; extra == "dev"
+Requires-Dist: opentelemetry-api~=1.33.1; extra == "dev"
 Requires-Dist: setuptools<80.0.0; extra == "dev"
-Requires-Dist: nvidia-modelopt[torch]~=0.31.0; sys_platform != "darwin" and extra == "dev"
-Requires-Dist: megatron-energon[av_decode]<7; extra == "dev"
+Requires-Dist: nvidia-modelopt[torch]<0.32.0,>=0.31.0a0; sys_platform != "darwin" and extra == "dev"
+Requires-Dist: megatron-energon[av_decode]~=6.0; extra == "dev"
+Requires-Dist: flashinfer-python; extra == "dev"
 Provides-Extra: lts
 Requires-Dist: tqdm; extra == "lts"
 Requires-Dist: einops; extra == "lts"

{megatron_core-0.14.0rc2 → megatron_core-0.14.0rc3}/megatron/core/distributed/custom_fsdp/fully_sharded_data_parallel.py

@@ -217,6 +217,16 @@ class FullyShardedDataParallel(_BaseDataParallel):
 
         self.module.apply(unmap_weight_tensor)
 
+        for param in self.module.parameters():
+            if not hasattr(param, 'grad_added_to_main_grad'):
+                # This is to ensure that the param.grad_added_to_main_grad is set to False
+                # when the parameter is created.
+                param.grad_added_to_main_grad = False
+            if not hasattr(param, '__fsdp_param__'):
+                # This is to ensure that the param.__fsdp_param__ is set to True
+                # when the parameter is created.
+                param.__fsdp_param__ = True
+
     def _init_fsdp_param_and_grad_buffer(self):
         if self.config.calculate_per_token_loss:
             # We don't need to scale the gradients in this case.

{megatron_core-0.14.0rc2 → megatron_core-0.14.0rc3}/megatron/core/fp8_utils.py

@@ -2,7 +2,9 @@
 
 """Utility functions related to FP8 that are used throughout Megatron core"""
 
+import weakref
 from contextlib import nullcontext
+from functools import wraps
 from typing import List, Optional
 
 import torch
@@ -53,6 +55,29 @@ except (ImportError, ModuleNotFoundError):
     # MXFP8Tensor not found
     HAVE_TE_MXFP8TENSOR = False
 
+if HAVE_TE:
+    from megatron.core.extensions.transformer_engine import (
+        TEColumnParallelLinear,
+        TELayerNormColumnParallelLinear,
+        TELinear,
+        TERowParallelLinear,
+    )
+
+    TE_LINEAR_TYPES = (
+        TELinear,
+        TEColumnParallelLinear,
+        TERowParallelLinear,
+        TELayerNormColumnParallelLinear,
+    )
+else:
+    TE_LINEAR_TYPES = ()
+
+try:
+    from megatron.core.extensions.transformer_engine import Fp8Padding, Fp8Unpadding
+except ImportError:
+    Fp8Padding = None
+    Fp8Unpadding = None
+
 
 def is_float8tensor(tensor: torch.Tensor) -> bool:
     """Check if a tensor is a Transformer Engine Float8Tensor.
@@ -511,3 +536,97 @@ else:
     def get_fp8_context(config: TransformerConfig, layer_no: int = -1, is_init: bool = False):
         """Returns dummy fp8 context manager since TE is not available."""
         return nullcontext()
+
+
+if HAVE_TE:
+    from transformer_engine.pytorch.fp8 import FP8GlobalStateManager
+
+    # Modules that have been wrapped for inference for fp8
+    _fp8_inference_wrapped_modules = weakref.WeakSet()
+
+    def _wrap_te_linear_for_padding(module: torch.nn.Module):
+        """Wrap a TE linear module to automatically pad sequences for FP8 inference.
+
+        Modifies the module's forward method to:
+        1. Pad input sequences to FP8 alignment requirements
+        2. Run the original forward pass
+        3. Unpad outputs to original sequence length
+
+        Args:
+            module: A Transformer Engine linear layer (TELinear, TEColumnParallelLinear, etc.)
+        """
+        if module in _fp8_inference_wrapped_modules:
+            return
+        _pad_func = Fp8Padding(1)
+        _unpad_func = Fp8Unpadding(1)
+
+        original_forward = module.forward
+
+        @wraps(original_forward)
+        def padded_forward(input_tensor, *args, **kwargs):
+            # Only do padding for fp8 if we are in fp8 context
+            if not FP8GlobalStateManager.is_fp8_enabled():
+                return original_forward(input_tensor, *args, **kwargs)
+
+            seq_len, batch_size, hidden_size = input_tensor.shape
+            # Reshape to (S, B*H) to pad sequence dimension
+            input_2d = input_tensor.reshape(seq_len, -1)
+            # Pad the sequence dimension
+            padded_input_2d, _ = _pad_func(input_2d, [seq_len])
+            padded_seq_len = padded_input_2d.shape[0]
+
+            # Reshape back to (padded_S, B, H)
+            padded_input_3d = padded_input_2d.view(padded_seq_len, batch_size, hidden_size)
+            output = original_forward(padded_input_3d, *args, **kwargs)
+
+            # Handle output
+            if isinstance(output, tuple):
+                output_tensor = output[0]
+                other_outputs = output[1:]
+            else:
+                output_tensor = output
+                other_outputs = ()
+
+            # Unpad output - reshape to 2D, unpad, reshape back
+            _, _, output_hidden_size = output_tensor.shape
+            output_2d = output_tensor.reshape(padded_seq_len, -1)
+            unpadded_output_2d = _unpad_func(output_2d, [seq_len])
+            unpadded_output = unpadded_output_2d.reshape(seq_len, batch_size, output_hidden_size)
+
+            if other_outputs:
+                return (unpadded_output,) + other_outputs
+            else:
+                return unpadded_output
+
+        module.forward = padded_forward
+        _fp8_inference_wrapped_modules.add(module)
+
+    def prepare_model_for_fp8_inference(model):
+        """Prepare a model for FP8 inference by wrapping TE linear layers with padding support.
+
+        FP8 TE Gemms have specific shape requirements. This function wraps all Transformer
+        Engine linear layers in the model to automatically pad/unpad sequences during inference.
+
+        Args:
+            model (model (GPTModel): Model containing TE linear layers.
+
+        Returns:
+            GPTModel: The same model with wrapped linear layers (modified in-place).
+
+        """
+        assert Fp8Padding and Fp8Unpadding, "TE version does not have FP8 padding functions"
+        # Find and wrap all TE linear layers
+        for module in model.modules():
+            if isinstance(module, TE_LINEAR_TYPES):
+                _wrap_te_linear_for_padding(module)
+
+        return model
+
+else:
+
+    def prepare_model_for_fp8_inference(model):
+        """If trys using prepare_model_for_fp8_inference without TE we error"""
+        raise RuntimeError(
+            "prepare_model_for_fp8_inference requires Transformer Engine to be installed. "
+            "Please install transformer-engine to use FP8 inference."
+        )

{megatron_core-0.14.0rc2 → megatron_core-0.14.0rc3}/megatron/core/inference/contexts/__init__.py

@@ -14,6 +14,7 @@ warnings.warn(
     DeprecationWarning,
 )
 from .dynamic_context import (
+    ActiveRequestCountOverflowError,
     ChunkOverflowError,
     ContextOverflowError,
     DynamicInferenceContext,

{megatron_core-0.14.0rc2 → megatron_core-0.14.0rc3}/megatron/core/inference/contexts/dynamic_context.py

@@ -56,6 +56,18 @@ class ChunkOverflowError(ContextOverflowError):
     pass
 
 
+class ActiveRequestCountOverflowError(ContextOverflowError):
+    '''Used when `initialize_attention_state()` is called with
+    `num_warmup_requests > max_requests.'''
+
+    def __init__(self, max_request_count, active_request_count):
+        assert active_request_count > max_request_count
+        super().__init__(
+            "active_request_count (%d) > max_request_count (%d)."
+            % (active_request_count, max_request_count)
+        )
+
+
 # pylint: disable=line-too-long
 class DynamicInferenceContext(BaseInferenceContext):
     """Inference context that is passed to the main model in order
@@ -108,6 +120,11 @@ class DynamicInferenceContext(BaseInferenceContext):
             from `buffer_overflow_factor`.
         max_tokens_override (Optional[int]): If set, overrides value computed
             from `buffer_overflow_factor`.
+        tensor_model_parallel_size (Optional[int]): Tensor model parallel size.
+        num_cuda_graphs (Optional[int]): Maximum number of cuda graphs to capture,
+            where the cuda graph batch sizes range from 1 to `max_requests` (as
+            computed below). Due to rounding, the actual number of cuda graphs may
+            not equal this argument.
     """
 
     def __init__(
@@ -125,6 +142,7 @@ class DynamicInferenceContext(BaseInferenceContext):
         max_requests_override: Optional[int] = None,
         max_tokens_override: Optional[int] = None,
         tensor_model_parallel_size: Optional[int] = None,
+        num_cuda_graphs: Optional[int] = None,
         materialize_only_last_token_logits: bool = True,
     ):
 
@@ -188,7 +206,7 @@ class DynamicInferenceContext(BaseInferenceContext):
         self.active_token_count = 0
         self.paused_request_count = 0
         self.padded_active_token_count = None
-        self.padded_active_sample_count = None
+        self.padded_active_request_count = None
         self.paused_tokens = None
 
         # Per-request state.
@@ -246,6 +264,34 @@ class DynamicInferenceContext(BaseInferenceContext):
             device=torch.cuda.current_device(),
         )
 
+        # Cuda graph request counts (i.e., batch sizes used for decode-only steps).
+        self.cuda_graph_request_counts = None
+        if num_cuda_graphs is not None:
+
+            # Ensure valid num_cuda_graphs.
+            num_cuda_graphs = min(max(num_cuda_graphs, 1), self.max_requests)
+
+            # Cuda graph step size.
+            cuda_graph_rounder = 8
+            self.cuda_graph_step_size = self.max_requests / num_cuda_graphs
+            self.cuda_graph_step_size = cuda_graph_rounder * int(
+                math.ceil(int(self.cuda_graph_step_size) / cuda_graph_rounder)
+            )
+
+            # Cuda graph request counts.
+            if num_cuda_graphs == 1:
+                self.cuda_graph_request_counts = [self.max_requests]
+            else:
+                self.cuda_graph_request_counts = list(
+                    range(self.cuda_graph_step_size, self.max_requests, self.cuda_graph_step_size)
+                )
+                if self.cuda_graph_request_counts[-1] != self.max_requests:
+                    self.cuda_graph_request_counts.append(self.max_requests)
+                self.cuda_graph_request_counts.reverse()
+
+            # Set used for validating active cuda graph request count.
+            self.cuda_graph_request_counts_set = set(self.cuda_graph_request_counts)
+
         # `*_decode_only` tensors are for use with cuda graphs to maintain
         # consistent input shapes, which is required to use cuda graphs. Cuda
         # graphs are used only during decode-only steps (i.e., no requests are in
@@ -269,7 +315,7 @@ class DynamicInferenceContext(BaseInferenceContext):
             (self.max_requests + 1,), 0, dtype=torch.int32, device=torch.cuda.current_device()
         )
 
-        self.kv_memory_decode_only = torch.full(
+        self.request_to_kv_chunk_ids_decode_only = torch.full(
             (self.max_requests, self.max_kv_chunk_count),
             0,
             dtype=torch.int,
@@ -278,27 +324,22 @@ class DynamicInferenceContext(BaseInferenceContext):
 
         # Guaranteed active requests.
         # * See details in the class docstring above. `gtd_request_fraction` is
-        #   the fraction of the memory buffer that is reserved for guaranteeing
-        #   that some number of active requests can always proceed with their
-        #   generations. The number of bytes defined by `gtd_request_fraction *
-        #   buffer_size_gb` is converted to a number of requests that this
-        #   reserved space can handle (`gtd_request_count`), and rounded to be an
-        #   exact multiple of `max_sequence_length`. This is then converted into
-        #   the number of reserved chunks (`gtd_chunk_count`) and bytes
-        #   (`gtd_byte_count`).
-        # Chunk ids.
-        self.max_kv_chunk_count = math.ceil(self.max_sequence_length / self.chunk_size_tokens)
-        gtd_byte_count = buffer_guaranteed_fraction * buffer_size_bytes
-        gtd_request_count, _ = bytes_to_max_requests_and_tokens(gtd_byte_count)
-        if buffer_guaranteed_fraction > 0:
-            gtd_request_count = max(1, gtd_request_count)
-        gtd_request_count = self.round_up_requests(min(gtd_request_count, self.max_requests))
-        gtd_chunk_count = gtd_request_count * self.max_kv_chunk_count
-        assert (
-            gtd_request_count <= self.max_requests
-        ), "gtd_request_count (%d) > max_requests (%d)." % (gtd_request_count, self.max_requests)
-        self.gtd_request_count = gtd_request_count
-        self.gtd_chunk_count = gtd_chunk_count
+        #   the fraction of chunks in the memory buffer that are reserved for
+        #   guaranteeing that some number of active requests can always proceed
+        #   with their generations. The number of chunks defined by
+        #   `buffer_guaranteed_fraction * chunk_count_total` is converted to a
+        #   number of requests that this reserved space can safely handle
+        #   (`gtd_request_count`).
+        # * Note: computing the size of this guaranteed space from chunks rather
+        #   than bytes is safer due to the non-linear impacts of a large
+        #   `chunk_size_tokens` or `max_kv_chunk_count`. When computing from
+        #   chunks, this space will always be less than `chunk_count_total`. When
+        #   computing from bytes, this space can unexpectedly be much larger than
+        #   `chunk_count_total`, resulting in stalled generations.
+        gtd_chunk_count = int(buffer_guaranteed_fraction * chunk_count_total)
+        gtd_chunk_count = min(gtd_chunk_count, chunk_count_total)
+        self.gtd_request_count = max(1, gtd_chunk_count // self.max_kv_chunk_count)
+        self.gtd_chunk_count = self.gtd_request_count * self.max_kv_chunk_count
 
         # Initialize chunk allocator
         self.chunk_allocator = ChunkAllocator(
@@ -368,12 +409,7 @@ class DynamicInferenceContext(BaseInferenceContext):
 
     def cu_kv_lengths(self) -> Tensor:
         """Cumulative key/value sequence lengths."""
-        return (
-            self.cu_kv_seq_lengths,
-            self.kv_seq_lengths,
-            self.kv_seq_lengths_decode_only,
-            self.max_seqlen_k,
-        )
+        return (self.cu_kv_seq_lengths, self.kv_seq_lengths, self.max_seqlen_k)
 
     def get_active_sequence_lengths(self) -> Tensor:
         """Total sequence length (query + key) for active requests."""
@@ -487,7 +523,7 @@ class DynamicInferenceContext(BaseInferenceContext):
         key_seq_idx = self.token_to_position_in_request[:n]
         key_emb = key_emb[key_seq_idx]
         if self.is_decode_only():
-            assert key.shape[0] == n == self.max_requests
+            assert key.shape[0] == n
             key = apply_rotary_pos_emb(
                 t=key[:n], freqs=key_emb[:n], config=config, cp_group=cp_group
             )
@@ -506,23 +542,65 @@ class DynamicInferenceContext(BaseInferenceContext):
         self.query_seq_lengths_decode_only.fill_(0)
         self.cu_kv_seq_lengths = None
         self.cu_kv_seq_lengths_decode_only.fill_(0)
+        self.kv_seq_lengths = None
         self.kv_seq_lengths_decode_only.fill_(0)
-        self.kv_memory_decode_only.fill_(0)
+        self.request_to_kv_chunk_ids_decode_only.fill_(0)
         self.block_table = None
 
-    def initialize_attention_state(self) -> None:
-        """Initialize attention state so that every layer can use it"""
+    def initialize_attention_state(self, *, num_warmup_requests: Optional[int] = None) -> None:
+        """Initialize attention state so that every layer can use it.
+
+        Args:
+            num_warmup_requests (Optional[int]): Number of requests to use for
+                warming up cuda graphs. Must be less than or equal to
+                `max_requests`.
+
+        Return:
+            None.
+        """
 
+        # Use of num_warmup_requests only for decode-only.
+        if num_warmup_requests is not None:
+            assert self.is_decode_only(), "cuda graph warmup requires decode-only mode."
+
+        # Active request count.
+        active_request_count = (
+            self.total_request_count - self.paused_request_count
+            if num_warmup_requests is None
+            else num_warmup_requests
+        )
+
+        # Active cuda graph count (if decode-only).
+        active_cuda_graph_request_count = None
+        if self.is_decode_only():
+            if active_request_count > self.max_requests:
+                raise ActiveRequestCountOverflowError(self.max_requests, active_request_count)
+
+            if self.cuda_graph_request_counts:
+                active_cuda_graph_request_count = (
+                    math.ceil(active_request_count / self.cuda_graph_step_size)
+                    * self.cuda_graph_step_size
+                )
+                active_cuda_graph_request_count = min(
+                    active_cuda_graph_request_count, self.max_requests
+                )
+                assert active_cuda_graph_request_count in self.cuda_graph_request_counts_set
+            else:
+                active_cuda_graph_request_count = self.max_requests
+
+        # Padded active token/request counts.
         self.padded_active_token_count = (
-            self.max_requests
+            active_cuda_graph_request_count
             if self.is_decode_only()
             else self.round_up_tokens(self.active_token_count)
         )
-        self.padded_active_sample_count = (
-            self.max_requests
+        self.padded_active_request_count = (
+            active_cuda_graph_request_count
             if self.is_decode_only()
             else (self.total_request_count - self.paused_request_count)
         )
+
+        # Update token position indexes.
         self.token_to_chunk_idx[self.active_token_count : self.padded_active_token_count] = (
             self.dummy_chunk_idx
         )
@@ -533,6 +611,7 @@ class DynamicInferenceContext(BaseInferenceContext):
             self.active_token_count : self.padded_active_token_count
         ] = 0
 
+        # Update cu_query_seq_lengths, max_seqlen_q.
         query_lengths = self.request_query_lengths[
             self.paused_request_count : self.total_request_count
         ]
@@ -540,9 +619,7 @@ class DynamicInferenceContext(BaseInferenceContext):
             self.query_seq_lengths_decode_only[
                 0 : self.total_request_count - self.paused_request_count
             ] = query_lengths
-            cu_query_lengths_decode_only = torch.cumsum(self.query_seq_lengths_decode_only, dim=0)
-            self.cu_query_seq_lengths_decode_only[1:] = cu_query_lengths_decode_only
-            self.cu_query_seq_lengths = self.cu_query_seq_lengths_decode_only
+            self.cu_query_seq_lengths = None  # ensure no accidental use
             self.max_seqlen_q = 1
         else:
             cu_query_lengths = torch.cumsum(query_lengths, dim=0)
@@ -558,12 +635,18 @@ class DynamicInferenceContext(BaseInferenceContext):
         kv_seq_lengths = self.request_kv_length_offsets + self.request_query_lengths
         self.kv_seq_lengths = kv_seq_lengths[self.paused_request_count : self.total_request_count]
         if self.is_decode_only():
+            # Re-assign `kv_seq_lengths` to be a view of the first
+            # `active_cuda_graph_request_count` tokens of `kv_seq_lengths_decode_only`,
+            # such that `kv_seq_lengths` has a static memory address and is therefore
+            # cuda graph compatible. This allows `kv_seq_lengths` to transition between,
+            # cuda graph sizes, which makes multi-batch-size cuda graphs possible.
             self.kv_seq_lengths_decode_only[
                 0 : self.total_request_count - self.paused_request_count
             ] = self.kv_seq_lengths
-            cu_kv_lengths_decode_only = torch.cumsum(self.kv_seq_lengths_decode_only, dim=0)
-            self.cu_kv_seq_lengths_decode_only[1:] = cu_kv_lengths_decode_only
-            self.cu_kv_seq_lengths = self.cu_kv_seq_lengths_decode_only
+            self.kv_seq_lengths = self.kv_seq_lengths_decode_only[
+                : self.padded_active_request_count
+            ]
+            self.cu_kv_seq_lengths = None  # ensure no accidental use
             self.max_seqlen_k = self.max_sequence_length
         else:
             self.cu_kv_seq_lengths = torch.full(
@@ -575,14 +658,17 @@ class DynamicInferenceContext(BaseInferenceContext):
             self.cu_kv_seq_lengths[1:] = torch.cumsum(self.kv_seq_lengths, dim=0)
             self.max_seqlen_k = self.kv_seq_lengths.max().item()
 
-        kv_memory = self.request_to_kv_chunk_ids[
+        # Update KV chunk IDs, block table.
+        request_to_kv_chunk_ids = self.request_to_kv_chunk_ids[
             self.paused_request_count : self.total_request_count
         ]
         if self.is_decode_only():
-            self.kv_memory_decode_only[0 : self.total_request_count - self.paused_request_count] = (
-                kv_memory
-            )
-            self.block_table = self.kv_memory_decode_only
+            self.request_to_kv_chunk_ids_decode_only[
+                0 : self.total_request_count - self.paused_request_count
+            ] = request_to_kv_chunk_ids
+            self.block_table = self.request_to_kv_chunk_ids_decode_only[
+                : self.padded_active_request_count
+            ]
         else:
             self.block_table = self.request_to_kv_chunk_ids[
                 self.paused_request_count : self.total_request_count
@@ -606,7 +692,7 @@ class DynamicInferenceContext(BaseInferenceContext):
         self.active_token_count = 0
         self.paused_request_count = 0
         self.padded_active_token_count = 0
-        self.padded_active_sample_count = 0
+        self.padded_active_request_count = 0
         self.paused_tokens = None
 
         # Reset request indexes.
@@ -632,21 +718,24 @@ class DynamicInferenceContext(BaseInferenceContext):
         self.chunk_allocator.reset()
         self.request_to_kv_chunk_ids.fill_(-1)
 
-    def current_input_ids(self) -> Tensor:
-        """Flattened input IDs for forward pass.
-
-        Return:
-            (Tensor) Flattened active input IDs.
-        """
-        return self.token_to_input_ids[: self.padded_active_token_count].unsqueeze(0)
+    def current_input_and_position_ids(
+        self, *, num_warmup_tokens: Optional[int] = None
+    ) -> Tuple[Tensor, Tensor]:
+        """Flattened input and position IDs for forward pass.
 
-    def current_position_ids(self) -> Tensor:
-        """Flattened position IDs for forward pass.
+        Args:
+            num_warmup_tokens (Optional[int]): Number of tokens to return for
+                warming up cuda graphs. Must be less than or equal to
+                `max_tokens`.
 
         Return:
-            (Tensor) Flattened active position IDs.
+            (Tuple[Tensor, Tensor]) Flattened active input and position IDs.
         """
-        return self.token_to_pos_ids[: self.padded_active_token_count].unsqueeze(0)
+        num_tokens = num_warmup_tokens or self.padded_active_token_count
+        return (
+            self.token_to_input_ids[:num_tokens].unsqueeze(0),
+            self.token_to_pos_ids[:num_tokens].unsqueeze(0),
+        )
 
     def last_token_logits(self, logits: Tensor) -> Tensor:
         """Last tokens of logits.