Skip to content

feat: add Falcon Mamba training chat templates with generation markers#5723

Open
DagaBhai wants to merge 4 commits into
huggingface:mainfrom
DagaBhai:add-falconmamba-generation-marker-clean
Open

feat: add Falcon Mamba training chat templates with generation markers#5723
DagaBhai wants to merge 4 commits into
huggingface:mainfrom
DagaBhai:add-falconmamba-generation-marker-clean

Conversation

@DagaBhai

@DagaBhai DagaBhai commented May 7, 2026

Copy link
Copy Markdown
Contributor

What does this PR do?

contributing to #5471

Adds a training-compatible variant of the Falcon Mamba chat template with {% generation %} / {% endgeneration %} markers, enabling return_assistant_tokens_mask=True for assistant-only loss masking in SFT training.

-falconmamba.jinja: template from tiiuae/falcon-7b-instruct
-falconmamba_training.jinja: modified the falconmamba.jinja to turn the training chat template
-get_training_chat_template: registered falcon mamba
-tests/test_chat_template_utils.py: added param "trl-internal-testing/tiny-FalconMambaForCausalLM", id="falconmamba"

Before submitting

AI writing disclosure

We welcome the use of AI tools to help with contributions. For transparency and to help us improve our review process, please indicate the level of AI involvement in this PR.

  • No AI usage: the PR was written entirely by a human.
  • AI-assisted: some parts were suggested or improved by AI, but the PR was written and reviewed by a human.
  • AI-generated: the PR was mostly or fully generated by an AI tool.

Who can review?

@qgallouedec (opened #5471)


Note

Low Risk
Low risk: changes are additive and gated by exact chat-template identity checks, with a small chance of mismatching upstream templates and failing to apply the training patch.

Overview
Adds Falcon Mamba support to get_training_chat_template by introducing falconmamba.jinja and a training-patched falconmamba_training.jinja that wraps assistant output in {% generation %} / {% endgeneration %} for correct return_assistant_tokens_mask=True behavior.

Updates trl/chat_templates/README.md to document the new templates and extends TestGetTrainingChatTemplate coverage to include the tiny-FalconMambaForCausalLM tokenizer.

Reviewed by Cursor Bugbot for commit e3ae2f6. Bugbot is set up for automated code reviews on this repo. Configure here.

Comment thread trl/chat_templates/falconmamba_training.jinja
Comment thread trl/chat_template_utils.py Outdated
@HuggingFaceDocBuilderDev

Copy link
Copy Markdown

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@cursor cursor Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cursor Bugbot has reviewed your changes and found 1 potential issue.

Fix All in Cursor

❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.

Reviewed by Cursor Bugbot for commit e3ae2f6. Configure here.

', '
') }}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ '

Assistant:' }}{% endif %} No newline at end of file

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Wrong model template used for Falcon Mamba

High Severity

The template in falconmamba.jinja is actually from tiiuae/falcon-7b-instruct (original Falcon, FalconForCausalLM), not from tiiuae/falcon-mamba-7b-instruct (Falcon Mamba, FalconMambaForCausalLM). The real Falcon Mamba model uses a ChatML-style template (<|im_start|>/<|im_end|>), which is completely different from the \n\nUser:/\n\nAssistant: format stored here. The exact-match comparison in get_training_chat_template will never trigger for actual Falcon Mamba users, making this feature non-functional for its stated purpose.

Additional Locations (1)
Fix in Cursor Fix in Web

Reviewed by Cursor Bugbot for commit e3ae2f6. Configure here.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

arrf good point

@DagaBhai DagaBhai May 8, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes i was using falcon mamba but in the test chat template it was falcon instruct so i had to use this one

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants