Tensor-parallel SSM #333

jlamypoirier · 2025-07-21T21:17:52Z

✨ Description

Please provide a brief summary of the changes, relevant motivation, and context.
Include any related issue numbers or links to discussions, and explain why this change is necessary.

Closes #

🔍 Type of change

Select all that apply:

🐛 Bug fix (non-breaking change that addresses a specific issue)
🚀 New feature (non-breaking change that adds functionality)
⚠️ Breaking change (a change that could affect existing functionality)
📈 Performance improvement/optimization (improves speed, memory usage, or efficiency)
🛠️ Code refactor (non-functional changes that improve code readability, structure, etc.)
📦 Dependency bump (updates dependencies, including Dockerfile or package changes)
📝 Documentation change (updates documentation, including new content or typo fixes)
🔧 Infrastructure/Build change (affects build process, CI/CD, or dependencies)

jlamypoirier · 2025-07-30T15:44:39Z

tests/models/test_checkpoint.py

@@ -284,10 +284,15 @@ def test_load_pretrained(
 @pytest.mark.model_testing_group(ModelTestingGroup.convert)
 def test_huggingface_model(model_testing_config, get_convert_path):


This test isn't working for SSM because the HF wrapper can't find the external model. I would like to make it work so we have at least one correctness test. Any idea how to make it work? (@bigximik @oleksost @tscholak ?)

In my Vision+Hybrid-ssm PR, I updated the SSM conversion to copy the modeling files to the export directory https://github.com/ServiceNow/Fast-LLM/pull/332/files#diff-58be369d99e6722a68e734002686ae4afcfd423261e4d3d3b9d6aa552a6f2a14R729-R784
But this PR is far from being merged ...

I managed to add them here too, but the external models don't seem to be working

FAILED tests/models/test_checkpoint.py::test_huggingface_model[hybrid_mamba2]@dependency_group_2 - AttributeError: 'DynamicCache' object has no attribute 'has_previous_state' FAILED tests/models/test_checkpoint.py::test_huggingface_model[hybrid_discrete_mamba2]@dependency_group_3 - AttributeError: 'NoneType' object has no attribute 'ssm_states'

oleksost · 2025-08-19T16:33:37Z

fast_llm/layers/ssm/mamba2.py

@@ -27,23 +27,7 @@
 except (ImportError, RuntimeError):
    _causal_conv1d_available = False

-
-def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:


Why not keeping this function?

From, what I'm seeing there is absolutely no benefit over calling repeat_interleave directly. I tried to figure out why it's there, got two hypotheses:

expand is preferable over repeat because of the absence of copy. That's pointless because the copy is still done in reshape below. And there is a second copy on each usage (contiguous), so the function actually makes things slower...

repeat_interleave may involve cuda synchronization because it supports tensor inputs. But that's not supposed to happen, and the explicit output_size ensures it.

oleksost · 2025-08-19T16:59:54Z

fast_llm/layers/ssm/config.py

+        else:
+            head_dim = state
+
+        tensor_space.add_tensor_dim(head_groups := TensorDim(SSMDimNames.head_groups, num_head_groups, tensor))


These var names are somewhat confusing. Wouldn't this be clearer?

num_head_groups --> num_heads_per_group -- this is the number of heads in each group (e.g. div(self.d_xb, self.state_size), where head dim. is self.state_size)

TensorDim head_groups --> heads_per_group

group_heads --> head_groups --- this is number groups (like in GKV), so number of groups with num_head_groups heads in each group

The names are supposed to be correct, i.e. head_groups and num_ head_groups refers to the different head groups and the number of such groups, and group_heads refers to heads inside a group. I might have inverted them by mistake though, I'll double-check. (They should be right here.)

I'm not a huge fan of the term group_heads though, so I'm open to renaming. (heads_per_group? heads_in_group?) What do you think?

I see, makes sense, so the current names are correct. I guess for group_heads it would make sense to call it heads_per_group

fast_llm/layers/ssm/mamba2.py

oleksost · 2025-08-20T13:42:29Z

fast_llm/layers/ssm/mamba2.py

+            bias=config.add_bias_linear,
+            weight_init_method=init_kaiming_(self._config.d_inner),
+            sequence_parallel=self._sequence_parallel,
+            # TODO: lr_scale?


lr_scale=lr_scale?

Just making sure it's not intentionally absent. It's ok to add?

Same about the normalization layer?

yeh, should be added everywhere, I think we missed it before.

oleksost · 2025-08-20T17:47:59Z

fast_llm/layers/ssm/config.py

+        )
+        tensor_space.add_tensor_dim(
+            head_groups_and_state := CompositeTensorDim(
+                SSMDimNames.composite_head_groups_and_state, (head_groups, state)


SSMDimNames.composite_head_groups_and_state, (head_groups, head_dim) instead of SSMDimNames.composite_head_groups_and_state, (head_groups, state)? state and head_dim must not be the same (currently they are)

In mamba and mamba 2 head_dim and state are the same, so the two are equivalent. The distinction is for discrete mamba.

This ended up being too confusing though, so in the next PR the dimensions are moved directly in the model with model-specific names https://github.com/ServiceNow/Fast-LLM/pull/339/files#diff-646273ef4d1b740f2ea28a29b34ca91517f5ea91ee1530ede13acfc44d30f1b1R55

Moving the dimentions into the model would be much less confusing!

Would it make sense to move those in the model already in this pr?

I'd prefer to limit the amount of extra work since #339 is already way ahead of this PR, but I'll check if I can do it easily.

Done, wasn't too bad

oleksost · 2025-08-21T12:32:33Z

fast_llm/layers/ssm/mamba2.py

+
+        # inner_projection : (batch/local_sequence, local_sequence/batch, hidden)
+        #   -> (batch/sequence, sequence/batch, inner_projection)
+        inner_projection = self.in_proj(input_)


Not sure this is the best place to discuss this, but I would like to better understand how TP is working here.
My understanding of the shapes is case of TP=2:

input_ -- bs x seq/2 x d

inner_projection -- seq x bs x d -- so here the output is already gathered across the ranks?

all the way down to out_proj we operate on full seq. length

the out_proj splits the seq. length in 2 again

Is that correct? So we only apply TP to in_proj and out_proj (+ also dt_proj). So does it mean that inner_projection is identical across the ranks and all the ranks are performing redundant work?

inner_projection is tensor-parallel, ex. with size seq x bs x d/2. I'll change to (batch/sequence, sequence/batch, local_inner_projection) to clear things up.

Sorry, I meant in case of sequence_tensor_parallel=True, where we split over sequence length. Sequence_tensor_parallel should work, right? (my experiments show that loss and grad-norms are very similar in STP and non-TP, so I assume it works).

So in TP=2 case, each rank processes half of the inner_projection, that works fine.

In TP=2 and sequence_tensor_parallel=True each rank processes half of the sequence dim.? I just don't quite understand how it works in this case, since input_ is bs x seq/2 x d, but inner_projection is already seq x bs x d --- so does it mean everything afterin_proj operates on full sequence length (which makes sense since selective_scan_fn does not support STP)?

I'm not sure if that's what you mean, but in sequence-tensor-parallel we split the activations that aren't already tensor-parallel over the sequence (mainly hidden states), so that all (or almost all) of them are split in the TP direction in one way or another. For mamba (and attn) that means the input is sequence-parallel, then we (gather the full tensor)[https://github.com/ServiceNow/Fast-LLM/blob/main/fast_llm/functional/linear.py#L114] and run inner_projection with standard TP which leaves a TP output.

For mamba (and attn) that means the input is sequence-parallel, then we (gather the full tensor)[https://github.com/ServiceNow/Fast-LLM/blob/main/fast_llm/functional/linear.py#L114] and run inner_projection with standard TP which leaves a TP output.
So does it mean if the input to mamba is sequence-parallel (bs x seq./tp x d), then in_proj gathers it into bs x seq. x d, then doing the projection with standard TP means it should output bs x seq. x d/tp?

jlamypoirier added 9 commits July 21, 2025 17:17

TP mamba

82eed2b

TP mamba

4e310c7

fix

3cc4118

fix

9f7f75c

fixes

4054e04

fix

0014cc6

fixes

47ad548

fixes

6a074fa

Update external

d66651f

jlamypoirier mentioned this pull request Jul 24, 2025

SSM debugging and TP prequisites #335

Merged

jlamypoirier added 2 commits July 24, 2025 15:19

Merge branch 'main' into tp_mamba

5006328

Merge branch 'debug_mamba' into tp_mamba

13176bd

jlamypoirier changed the base branch from main to debug_mamba July 24, 2025 19:21

jlamypoirier added 15 commits July 24, 2025 17:34

Merge branch 'debug_mamba' into tp_mamba

73f591f

stuff

1feccc8

misc

e528b50

Merge branch 'debug_mamba' into tp_mamba

bb4dcd9

misc

c1b7f44

Merge branch 'debug_mamba' into tp_mamba

051bb07

misc

0a9ff25

Parallel discrete mamba 2

e7d9636

Mamba 2, misc

c14b764

doc

b605bd2

Merge branch 'debug_mamba' into tp_mamba

0a3e2a7

fixes

2e6d082

misc

b6c8613

Merge branch 'debug_mamba' into tp_mamba

acdfab1

Merge branch 'concatenated_dim' into tp_mamba

93e4c94

jlamypoirier changed the base branch from debug_mamba to concatenated_dim July 28, 2025 22:11

doc

c41efc2

jlamypoirier added 2 commits July 28, 2025 19:38

cleanup

0b8bd5d

Merge branch 'concatenated_dim' into tp_mamba

c0f1597

jlamypoirier marked this pull request as ready for review July 29, 2025 22:04

jlamypoirier commented Jul 30, 2025

View reviewed changes

fix

cef7c15

jlamypoirier requested review from RaymondLi0, oleksost, nitsanluke, tscholak and bigximik and removed request for RaymondLi0 July 31, 2025 19:10

jlamypoirier added 4 commits August 8, 2025 13:55

Merge branch 'concatenated_dim' into tp_mamba

defd6e0

fixes

8abf258

Merge branch 'concatenated_dim' into tp_mamba

fd3307d

stuff

0e2e124

oleksost reviewed Aug 19, 2025

View reviewed changes

oleksost reviewed Aug 20, 2025

View reviewed changes

oleksost reviewed Aug 21, 2025

View reviewed changes

jlamypoirier added 7 commits August 21, 2025 14:41

Pr comments

9a2a7a2

Cleanup

8c382a9

Cleanup

019e43d

Cleanup

3e0f3e5

fixes

1abdd19

fixes

7c24292

fixes

af2964b

		@@ -284,10 +284,15 @@ def test_load_pretrained(
		@pytest.mark.model_testing_group(ModelTestingGroup.convert)
		def test_huggingface_model(model_testing_config, get_convert_path):

Tensor-parallel SSM #333

Are you sure you want to change the base?

Tensor-parallel SSM #333

Uh oh!

Conversation

jlamypoirier commented Jul 21, 2025

✨ Description

🔍 Type of change

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

oleksost Aug 20, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

oleksost Aug 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

oleksost Aug 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

oleksost Aug 20, 2025 •

edited

Loading

oleksost Aug 21, 2025 •

edited

Loading

oleksost Aug 22, 2025 •

edited

Loading