feat: add Falcon Mamba training chat templates with generation markers#5723
feat: add Falcon Mamba training chat templates with generation markers#5723DagaBhai wants to merge 4 commits into
Conversation
|
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. |
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes and found 1 potential issue.
❌ 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 |
There was a problem hiding this comment.
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)
Reviewed by Cursor Bugbot for commit e3ae2f6. Configure here.
There was a problem hiding this comment.
Yes i was using falcon mamba but in the test chat template it was falcon instruct so i had to use this one


What does this PR do?
contributing to #5471
Adds a training-compatible variant of the Falcon Mamba chat template with
{% generation %} / {% endgeneration %}markers, enablingreturn_assistant_tokens_mask=Truefor assistant-only loss masking in SFT training.-
falconmamba.jinja: template fromtiiuae/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
{% generation %}chat templates for common model families #5471AI 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.
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_templateby introducingfalconmamba.jinjaand a training-patchedfalconmamba_training.jinjathat wraps assistant output in{% generation %}/{% endgeneration %}for correctreturn_assistant_tokens_mask=Truebehavior.Updates
trl/chat_templates/README.mdto document the new templates and extendsTestGetTrainingChatTemplatecoverage to include thetiny-FalconMambaForCausalLMtokenizer.Reviewed by Cursor Bugbot for commit e3ae2f6. Bugbot is set up for automated code reviews on this repo. Configure here.