[feat] Hybrid Mamba model with Mamba and discrete Mamba 2 layers

[oleksost]

✨ Description

This PR integrates Mamba1 and discrete Mamba2 blocks into fast-llm training pypeline, this is the initial step to address #68 .
It introduces a basic hybrid architecture that can interleave transformer and mamba-1 blocks.

Next steps:

• add discrete mamba-2 from Llamba
• be able to load and export Llamba-8B to support Llamba support #197

The training with a simple hybrid model can be tested:

1. Install mamba_ssm and 'causal-conv1d' dependency, pip install mamba_ssm[causal-conv1d]==2.2.4
2. 2. launch training by passing

       "args": [
                   "train",
                   "hybrid_ssm",
                   "--config",
                   "path/to/hybrid_config.yaml"
               ],
   

and the following simple config to build a hybrid model:

model:
  base_model:
    transformer:
      num_layers: 6
      use_flash_attention: no  
    ssm:
      dt_rank: auto
      state_size: 16
      expansion_factor: 2
      debug_ssm: false
    block_pattern: ["m", "t", "m", "m2", "m", "m"] # mixing transformer, mamba 1 and descrete mamba layers
  
  distributed:
    training_dtype: bf16
    tensor_parallel: 1 
    pipeline_parallel: 1
    world_size: 1 

training:
  train_iters: 1000  
  logs:
    interval: 10
  validation:
    iterations: 25
    interval: 1000
  wandb:  
    project_name: fast-llm-ssm-test
    group_name: ssm
    entity_name: null

data:
  datasets:
    Training:
      type: memmap
      path: /home/toolkit/dev/fast-llm-tutorial/dataset/shard_0_0
    Validation:
      type: memmap
      path: /home/toolkit/dev/fast-llm-tutorial/dataset/shard_0_0

To load Llamba1B model, add the following to the config:

pretrained:
  format: llamba
  path: /mnt/checkpoints/pretrained_models/Llamba-1B

🔍 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)

📝 Changes

• added minimal mamba-1 layer and block
• added hybrid model and corresponding configs
• the implementation follows the one from https://github.com/Zyphra/Zamba2 and https://github.com/state-spaces/mamba

✅ Checklist

Make sure the following tasks are completed before submitting the PR:

General

• 📜 I have read and followed the contributing guidelines.
• 🏷️ I am using a clear and descriptive PR title that summarizes the key change or feature introduced.
• 🎉 The functionality is complete, and I have tested the changes.
• 📝 I have updated the documentation if needed.
• ⚠️ The change does not introduce any new issues (e.g., runtime warnings, type checker errors, linting problems, unhandled edge cases).
• 🧩 I have commented my code, especially in hard-to-understand areas.

Dependencies and Configuration

• 🐋 I have updated the Docker configuration or dependencies, if applicable.
• 🔄 I have ensured compatibility with the existing setup after dependency changes.

Testing

• 🧪 I have added or updated tests to cover my changes.
• ✔️ New and existing tests pass locally with my changes.
• 🚦 I have tested these changes on GPUs and verified training stability.
• 🏋️ I have tested the changes on realistic training workloads, if applicable.

Performance Impact

• 📊 I have run benchmarks where applicable to evaluate the performance impact.
• ✅ The benchmarks show no performance regression.
• 🚀 The benchmarks indicate a potential performance improvement.
• ⚠️ The benchmarks indicate a potential performance degradation.
• 📈 I have provided benchmark results and detailed any performance impact below, if applicable.

---

🗒️ Additional Notes

• currently some parameters that are used for defining hybrid model's architecture are in the transformer config, e.g. num_layers, but they should probably moved to higher level configs at some point

[jlamypoirier]

Redundant with block_pattern.

[oleksost]

Currently, this is needed to load Llamba model: we set default_block to m2 in _create_config_converters of LLambaHuggingfaceCheckpointHandler and the block_pattern is then created in the __post_init__ of HybridSSMBaseModelConfig. This is a bit cumbersome indeed.

[jlamypoirier]

Why can't it set block_pattern? Also this would need to go in _validate, not __post_init__

[oleksost]

afaiu _create_config_converters does not know about the num_layers in the loaded config, so we cannot set the block pattern as it depends on the number of layers.

Moved the post_init logic for block config to validate

[jlamypoirier]

Oh so block_pattern is not specifying a repeated pattern, but the entire list? Why not just repeating the list up to num_layers instead, as the name suggests?

[oleksost]

Renamed block_pattern into hybrid_block_layout and use default_block repeated 'num_layers' times in case 'hybrid_block_layout' is not specified.

[jlamypoirier]

Ok, but why not a single variable with automated repetition?

[tscholak]

@oleksost, try merging these but use this default factory:

    hybrid_block_layout: list[str] = Field(
        default_factory=lambda: ['m'],
        desc="Pattern of blocks to use in the model. 't' for Transformer, 'm' for Mamba1, 'm2' for Descrete Mamba2.",
        hint=FieldHint.core,
    )

that avoids the mutable default trap.

[jlamypoirier]

Oh so block_pattern is not specifying a repeated pattern, but the entire list? Why not just repeating the list up to num_layers instead, as the name suggests?

[jlamypoirier]

Please adjust field names for our naming conventions.

[jlamypoirier]

Please use None for derived defaults. dt_rank: int = Field(default=None, ...

[jlamypoirier]

Why would it be a dict? There must be another problem elsewhere.

[tscholak]

do we have tests that check if serializing/deserializing lists with strings works correctly? the special property of strings is that they are also lists, kind-of. there could be issues because of that, and they may be not captured by tests if we only test integers of ints, say.

[tscholak]

@oleksost I see a patch to set_nested_dict_value below. My hunch is that this isn't working properly

[jlamypoirier]

Not needed, you're already passing layer_idx to the mixer_cls() call

[jlamypoirier]

This doesn't work with TP, need to explicitly prevent. (Not sure about PP).

[tscholak]

we're gonna have TP eventually, not in scope for this PR though

[tscholak]

@jlamypoirier, naming-wise, "layer" implies a single functional unit, like attention or an MLP. But here, the repeated unit is a composite: It includes a mixer (attention, SSM, etc.), normalization, residuals, and post-mixer processing (MLPs, MoEs, etc.). This structure isn't atomic. It's multiple layers stitched into a reusable computation unit, which is typically referred to as a "block" in other model families (e.g., ResNet, Swin, and even “Transformer blocks” in many papers and blogs).
Now that we're supporting architectures beyond transformers (Mamba, etc.), the term “block” avoids misleading assumptions:

1. "Layer" strongly implies a fixed layout rooted in the transformer design.
2. "Block" reflects the actual structure: a swappable, composite computation unit.

Keeping the name BaseBlock lets us consistently subclass it for Mamba, attention, etc., without implying that all of them are "layers" in the same architectural tradition. In this sense, "TransformerLayer" can be a Block instance, but not every Block is a TransformerLayer.

Also, and I find this the most important argument: internally and casually, we already call them blocks. Making the code match mental models reduces friction.

[jlamypoirier]

@tscholak Sounds reasonable, problem s these things are called "layers" everywhere else in Fast-LLM. Should we think about renaming these too?

[tscholak]

I know we need this line to be removed for the debugger to work. You could do this instead:

    except Exception:  # noqa
        if sys.gettrace():
            raise
        logger.critical(traceback.format_exc())
        sys.exit(1)

[jlamypoirier]

That would work, we do need that line outside a debugger. Same thing needed for ValidationError above.

[tscholak]

This whole function is mysterious. It blindly creates empty dictionaries (setdefault(key, {})), even when we might actually want lists.
On top of this, do we need this patch? What does this do? The special case for lists is confusing. Do we have tests for the added behavior?

[jlamypoirier]

This method can't really support list, let's not try to make up a new handling rule for them.
I don't really see what would be the use for it here anyway, since it's easy to just pass a list of strings.
Also set_nested_dict_value has been rewritten in main...

[tscholak]

Looks already very good to me, thanks @oleksost
This should go in asap!
I had a few comments and suggestions. Mostly, I think we don't want to be too picky with this PR at this point because we're going to be actively working on improving many parts of this anyway in the next weeks.

[jlamypoirier]

Why the renaming?

[jlamypoirier]

I guess we can ignore some issues with the current SSM implementation, if we handle it properly. This means clear warnings that the model is experimental and may break at any point, ex. at model runtime and/or in file headers. We still need to fix code changes outside the model though (transformers.py and set_nested_dict_value)

Also keep in mind that future modifications may break experiment configs, pretrained models and checkpoints, hence the importance of getting good config and parameter structures as soon as possible.

[jlamypoirier]

That would work, we do need that line outside a debugger. Same thing needed for ValidationError above.

[jlamypoirier]

Import

[tscholak]

This is a good moment to reflect on the pattern.

Right now, Fast-LLM expects contributors to internalize a set of nuanced import rules (see https://servicenow.github.io/Fast-LLM/contributing/style-guide/#imports) that go beyond what most Python projects require. That may have worked at some point, but it doesn't scale. New contributors can't memorize this, and even returning ones keep tripping over it.

If this style is important, it needs to be enforced and automatically fixable through linting or a pre-commit hook. If it can't be, we should let go of it. Patterns that can't be learned quickly or applied automatically create friction and slow down the team.

@jlamypoirier, could you file a ticket outlining what it would take to automate this rule? That's the only sustainable way forward.

[tscholak]

This means clear warnings that the model is experimental and may break at any point, ex. at model runtime and/or in file headers.

@oleksost, can you emit a warning in the logger when someone tries to instantiate the config for this model class? that should be enough for now.