Skip to content

refactor output tracing in mamba #44003

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
omkar-334 wants to merge 3 commits into from
Closed

refactor output tracing in mamba #44003

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
374a9e2
refactor output tracing in `mamba`
omkar-334 Feb 14, 2026
0921543
fix-repo
omkar-334 Feb 17, 2026
93a1439
Merge branch 'main' into refactor-mamba
omkar-334 Feb 17, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 9 additions & 31 deletions src/transformers/models/falcon_mamba/modeling_falcon_mamba.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,8 +33,9 @@
from ...integrations import lazy_load_kernel
from ...modeling_layers import GradientCheckpointingLayer
from ...modeling_utils import PreTrainedModel
from ...utils import ModelOutput, auto_docstring, logging
from ...utils import ModelOutput, auto_docstring, can_return_tuple, logging
from ...utils.import_utils import is_mambapy_available, is_torchdynamo_compiling
from ...utils.output_capturing import capture_outputs
from .configuration_falcon_mamba import FalconMambaConfig


Expand Down Expand Up @@ -557,6 +558,7 @@ class FalconMambaPreTrainedModel(PreTrainedModel):
_no_split_modules = ["FalconMambaBlock", "FalconMambaMixer"]
supports_gradient_checkpointing = True
_is_stateful = True
_can_record_outputs = {"hidden_states": FalconMambaBlock}

@torch.no_grad()
def _init_weights(self, module):
Expand Down Expand Up @@ -683,31 +685,26 @@ def get_input_embeddings(self):
def set_input_embeddings(self, new_embeddings):
self.embeddings = new_embeddings

@capture_outputs
@auto_docstring
def forward(
self,
input_ids: torch.LongTensor | None = None,
inputs_embeds: torch.LongTensor | None = None,
cache_params: FalconMambaCache | None = None,
use_cache: bool | None = None,
output_hidden_states: bool | None = None,
return_dict: bool | None = None,
cache_position: torch.LongTensor | None = None,
attention_mask: torch.LongTensor | None = None,
**kwargs,
) -> tuple | FalconMambaOutput:
) -> FalconMambaOutput:
r"""
cache_params (`FalconMambaCache`, *optional*):
If passed along, the model uses the previous state in all the blocks (which will give the output for the
`input_ids` provided as if the model add `state_input_ids + input_ids` as context).
use_cache (`bool`, *optional*):
If set to `True`, the `cache_params` is returned and can be used to quickly generate the next logits.
"""
output_hidden_states = (
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
)
use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False)
return_dict = return_dict if return_dict is not None else self.config.use_return_dict

if (input_ids is None) ^ (inputs_embeds is not None): # ^ is python for xor
raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
Expand Down Expand Up @@ -737,7 +734,6 @@ def forward(
cache_params = None

hidden_states = inputs_embeds
all_hidden_states = () if output_hidden_states else None
for mixer_block in self.layers:
hidden_states = mixer_block(
hidden_states,
Expand All @@ -746,21 +742,11 @@ def forward(
attention_mask=attention_mask,
)

if output_hidden_states:
all_hidden_states = all_hidden_states + (hidden_states,)

hidden_states = self.norm_f(hidden_states)

if output_hidden_states:
all_hidden_states = all_hidden_states + (hidden_states,)

if not return_dict:
return tuple(v for v in [hidden_states, cache_params, all_hidden_states] if v is not None)

return FalconMambaOutput(
last_hidden_state=hidden_states,
cache_params=cache_params if use_cache else None,
hidden_states=all_hidden_states,
)


Expand Down Expand Up @@ -846,6 +832,7 @@ def prepare_inputs_for_generation(

return model_inputs

@can_return_tuple
@auto_docstring
def forward(
self,
Expand All @@ -854,13 +841,11 @@ def forward(
inputs_embeds: torch.FloatTensor | None = None,
cache_params: FalconMambaCache | None = None,
labels: torch.LongTensor | None = None,
output_hidden_states: bool | None = None,
return_dict: bool | None = None,
use_cache: bool | None = None,
cache_position: torch.Tensor | None = None,
logits_to_keep: int | torch.Tensor = 0,
**kwargs, # for now we need this for generation
) -> tuple | FalconMambaCausalLMOutput:
) -> FalconMambaCausalLMOutput:
r"""
cache_params (`FalconMambaCache`, *optional*):
If passed along, the model uses the previous state in all the blocks (which will give the output for the
Expand All @@ -872,20 +857,17 @@ def forward(
use_cache (`bool`, *optional*):
If set to `True`, the `cache_params` is returned and can be used to quickly generate the next logits.
"""
return_dict = return_dict if return_dict is not None else self.config.use_return_dict

falcon_mamba_outputs = self.backbone(
input_ids,
cache_params=cache_params,
inputs_embeds=inputs_embeds,
output_hidden_states=output_hidden_states,
return_dict=return_dict,
use_cache=use_cache,
cache_position=cache_position,
attention_mask=attention_mask,
**kwargs,
)

hidden_states = falcon_mamba_outputs[0]
hidden_states = falcon_mamba_outputs.last_hidden_state
# Only compute necessary logits
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, :].to(self.lm_head.weight.dtype)).float()
Expand All @@ -901,10 +883,6 @@ def forward(
loss_fct = CrossEntropyLoss()
loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))

if not return_dict:
output = (logits,) + falcon_mamba_outputs[1:]
return ((loss,) + output) if loss is not None else output

return FalconMambaCausalLMOutput(
loss=loss,
logits=logits,
Expand Down
65 changes: 28 additions & 37 deletions src/transformers/models/mamba/modeling_mamba.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,9 +31,11 @@
from ...utils import (
ModelOutput,
auto_docstring,
can_return_tuple,
logging,
)
from ...utils.import_utils import is_mambapy_available, is_torchdynamo_compiling
from ...utils.output_capturing import capture_outputs
from .configuration_mamba import MambaConfig


Expand Down Expand Up @@ -216,7 +218,13 @@ def __init__(self, config: MambaConfig, layer_idx: int):

def warn_slow_implementation(self):
is_fast_path_available = all(
(selective_state_update, selective_scan_fn, causal_conv1d_fn, causal_conv1d_update, mamba_inner_fn)
(
selective_state_update,
selective_scan_fn,
causal_conv1d_fn,
causal_conv1d_update,
mamba_inner_fn,
)
)
if not is_fast_path_available:
if self.use_mambapy:
Expand Down Expand Up @@ -339,7 +347,7 @@ def cuda_kernels_forward(
return contextualized_states

# fmt: off
def slow_forward(self, input_states, cache_params: MambaCache | None=None, cache_position:torch.LongTensor | None=None, attention_mask: torch.LongTensor | None = None):
def slow_forward(self, input_states, cache_params: MambaCache | None = None, cache_position: torch.LongTensor | None = None, attention_mask: torch.LongTensor | None = None):
batch_size, seq_len, _ = input_states.shape
dtype = input_states.dtype
# 1. Gated MLP's linear projection
Expand Down Expand Up @@ -388,19 +396,19 @@ def slow_forward(self, input_states, cache_params: MambaCache | None=None, cache
ssm_parameters, [self.time_step_rank, self.ssm_state_size, self.ssm_state_size], dim=-1
)
discrete_time_step = self.dt_proj(time_step) # [batch, seq_len, intermediate_size]
discrete_time_step = nn.functional.softplus(discrete_time_step).transpose(1, 2) # [batch, intermediate_size, seq_len]
discrete_time_step = nn.functional.softplus(discrete_time_step).transpose(1, 2) # [batch, intermediate_size, seq_len]

# 3.b. Discretization: B and C to [batch, seq_len, intermediate_size, ssm_state_size] (SRAM)
A = -torch.exp(self.A_log.float()) # [intermediate_size, ssm_state_size]
discrete_A = torch.exp(A[None, :, None, :] * discrete_time_step[:, :, :, None]) # [batch, intermediate_size, seq_len, ssm_state_size]
discrete_A = torch.exp(A[None, :, None, :] * discrete_time_step[:, :, :, None]) # [batch, intermediate_size, seq_len, ssm_state_size]
discrete_B = discrete_time_step[:, :, :, None] * B[:, None, :, :].float() # [batch, intermediate_size, seq_len, ssm_state_size]
deltaB_u = discrete_B * hidden_states[:, :, :, None].float()

# 3.c perform the recurrence y ← SSM(A, B, C)(x)
if self.use_mambapy and self.training and cache_params is None:
hs = pscan(discrete_A.transpose(1, 2), deltaB_u.transpose(1, 2)) # [batch, seq_len, intermediate_size, ssm_state_size]
hs = pscan(discrete_A.transpose(1, 2), deltaB_u.transpose(1, 2)) # [batch, seq_len, intermediate_size, ssm_state_size]

scan_output = (hs @ C.unsqueeze(-1)).squeeze(3).transpose(1, 2) # [batch, intermediate_size, seq_len]
scan_output = (hs @ C.unsqueeze(-1)).squeeze(3).transpose(1, 2) # [batch, intermediate_size, seq_len]
scan_output = scan_output + hidden_states * self.D[None, :, None]
scan_output = scan_output * self.act(gate)
else:
Expand Down Expand Up @@ -429,7 +437,13 @@ def forward(
attention_mask: torch.LongTensor | None = None,
):
is_fast_path_available = all(
(selective_state_update, selective_scan_fn, causal_conv1d_fn, causal_conv1d_update, mamba_inner_fn)
(
selective_state_update,
selective_scan_fn,
causal_conv1d_fn,
causal_conv1d_update,
mamba_inner_fn,
)
)
if is_fast_path_available and "cuda" in self.x_proj.weight.device.type and not is_torchdynamo_compiling():
return self.cuda_kernels_forward(hidden_states, cache_params, cache_position, attention_mask)
Expand Down Expand Up @@ -491,6 +505,7 @@ class MambaPreTrainedModel(PreTrainedModel):
_no_split_modules = ["MambaBlock", "MambaMixer"]
supports_gradient_checkpointing = True
_is_stateful = True
_can_record_outputs = {"hidden_states": MambaBlock}

@torch.no_grad()
def _init_weights(self, module):
Expand Down Expand Up @@ -619,31 +634,26 @@ def get_input_embeddings(self):
def set_input_embeddings(self, new_embeddings):
self.embeddings = new_embeddings

@capture_outputs
@auto_docstring
def forward(
self,
input_ids: torch.LongTensor | None = None,
inputs_embeds: torch.LongTensor | None = None,
cache_params: MambaCache | None = None,
use_cache: bool | None = None,
output_hidden_states: bool | None = None,
return_dict: bool | None = None,
cache_position: torch.LongTensor | None = None,
attention_mask: torch.LongTensor | None = None,
**kwargs,
) -> tuple | MambaOutput:
) -> MambaOutput:
r"""
cache_params (`MambaCache`, *optional*):
If passed along, the model uses the previous state in all the blocks (which will give the output for the
`input_ids` provided as if the model add `state_input_ids + input_ids` as context).
use_cache (`bool`, *optional*):
If set to `True`, the `cache_params` is returned and can be used to quickly generate the next logits.
"""
output_hidden_states = (
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
)
use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False)
return_dict = return_dict if return_dict is not None else self.config.use_return_dict

if (input_ids is None) ^ (inputs_embeds is not None): # ^ is python for xor
raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
Expand Down Expand Up @@ -673,7 +683,6 @@ def forward(
cache_params = None

hidden_states = inputs_embeds
all_hidden_states = () if output_hidden_states else None
for mixer_block in self.layers:
hidden_states = mixer_block(
hidden_states,
Expand All @@ -682,21 +691,11 @@ def forward(
attention_mask=attention_mask,
)

if output_hidden_states:
all_hidden_states = all_hidden_states + (hidden_states,)

hidden_states = self.norm_f(hidden_states)

if output_hidden_states:
all_hidden_states = all_hidden_states + (hidden_states,)

if not return_dict:
return tuple(v for v in [hidden_states, cache_params, all_hidden_states] if v is not None)

return MambaOutput(
last_hidden_state=hidden_states,
cache_params=cache_params if use_cache else None,
hidden_states=all_hidden_states,
)


Expand Down Expand Up @@ -782,6 +781,7 @@ def prepare_inputs_for_generation(

return model_inputs

@can_return_tuple
@auto_docstring
def forward(
self,
Expand All @@ -790,13 +790,11 @@ def forward(
inputs_embeds: torch.FloatTensor | None = None,
cache_params: MambaCache | None = None,
labels: torch.LongTensor | None = None,
output_hidden_states: bool | None = None,
return_dict: bool | None = None,
use_cache: bool | None = None,
cache_position: torch.Tensor | None = None,
logits_to_keep: int | torch.Tensor = 0,
**kwargs, # for now we need this for generation
) -> tuple | MambaCausalLMOutput:
) -> MambaCausalLMOutput:
r"""
cache_params (`MambaCache`, *optional*):
If passed along, the model uses the previous state in all the blocks (which will give the output for the
Expand All @@ -808,20 +806,17 @@ def forward(
use_cache (`bool`, *optional*):
If set to `True`, the `cache_params` is returned and can be used to quickly generate the next logits.
"""
return_dict = return_dict if return_dict is not None else self.config.use_return_dict

mamba_outputs = self.backbone(
input_ids,
cache_params=cache_params,
inputs_embeds=inputs_embeds,
output_hidden_states=output_hidden_states,
return_dict=return_dict,
use_cache=use_cache,
cache_position=cache_position,
attention_mask=attention_mask,
**kwargs,
)

hidden_states = mamba_outputs[0]
hidden_states = mamba_outputs.last_hidden_state
# Only compute necessary logits
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, :].to(self.lm_head.weight.dtype)).float()
Expand All @@ -837,10 +832,6 @@ def forward(
loss_fct = CrossEntropyLoss()
loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))

if not return_dict:
output = (logits,) + mamba_outputs[1:]
return ((loss,) + output) if loss is not None else output

return MambaCausalLMOutput(
loss=loss,
logits=logits,
Expand Down