[rl] Add Muse Glimmer renderer and Search-R1 RL config - #4145
Conversation
meta-cla
Bot
added
the
CLA Signed
4da48c6 to
0c14084
Compare
|
Hey @HosseinKaviani-H , I tried to test the PR and ran into some issues: Two things blocked it from running on the public release (just want to make sure that you ran it against the HF checkpoints on not ckpnts on Manifold) and last one is more an env config missing, 1. Renderer config import crash: fields not classifiedOn import,
class MuseGlimmerRendererConfig(BaseRendererConfig):
name: Literal["muse_glimmer"] = RENDERER_NAME
# required by the current renderers lib: classify non-base fields
_internal_fields = frozenset(
{"reasoning_strength", "retain_reasoning_in_history", "answer_from_reasoning_fallback"}
)
...This looks like 2. (Main issue) Renderer depends on a Jinja
|
@AlirezaShamsoshoara Thanks for pointing these out. I fixed the issues 1 and 3. For 1, one change to your suggestion: reasoning_strength goes in _template_fields rather than _internal_fields, since it's forwarded to apply_chat_template verbatim and is exactly the kind of field the template-parity matrix is meant to cover. The other two are internal knobs with no Jinja analogue. For 2, I see your point. We want to adapt for the public release. I am running some experiments on this, let me get back to you. |
@HosseinKaviani-H awesome, and thanks. I will try again once we have a solution for 2. two suggestions by Claude:
|
| def rl_grpo_muse_glimmer_30b_search_r1() -> Controller.Config: | ||
| """GRPO/DAPO Search-R1 for Muse Glimmer 30B. | ||
|
|
||
| 8 GPUs: 6 trainer (FSDP=3 x TP=2) + 2 generator (TP=2), with a dense retrieval |
There was a problem hiding this comment.
FSDP=3 is very rare and do we really need 6 GPUs to fit the trainer? Does FSDP2 TP2 work?
There was a problem hiding this comment.
No. The model has 2 KV heads which caps TP at 2 for both trainer and generator. The generator takes 2 GPUs, leaving 6 and with TP=2 we would have FSDP=3. Model doesn't fit in 4 GPUs.
There was a problem hiding this comment.
Would this go to / be replaced by renderer side change?
There was a problem hiding this comment.
This one is more muse_glimmer native tho - structure is shared with gpt-oss, but the ATEM tool-call syntax is ours. If it goes upstream it'd be a tool parser in renderers/parsers.py next to the qwen3 / glm / deepseek-v3 ones. That's my understanding, but open to suggestions.
|
|
I don't think I understand this. From what I understand, "entropy" is measured across logits and it being high means the model doesn't know what to do. All 8 score the same and no gradient, wouldn't necessarily mean that the model gets more confused than before? |
You're right, my explanation was off. Zero-gradient groups do nothing, so it's the steps that did have gradient. Near the max the advantage gets uneven. The few failures get a much bigger push than the many successes, so the update is mostly pushing away from what the failures wrote. Pushing a sample down spreads its probability over every other token, while pushing a good one up does little since those tokens are already near certain. Clip-higher pushes the same way and there's no KL term holding it back. And you're right that it doesn't mean the model got more confused, accuracy didn't drop, only the variety of wording went up. If that's the cause, entropy should track the partially-correct groups rather than the fully-correct ones, so I'll log that next run. |
|
@yichuan-w Could I get your review for the renderer? Thanks! |
LGTM! |
77137b4 to
5a9288e
Compare
| @@ -73,6 +79,16 @@ def build(self, *, tokenizer_path: str) -> Renderer: | |||
|
|
|||
| # `name=None` (or "auto") -> let `create_renderer` resolve from the tokenizer. | |||
There was a problem hiding this comment.
@tianyu-l Can I get your review for this part? I have noticed there was a refactor and the renderer is now built in the spawned RolloutWorker, which never imports the config module. So register() there never ran and every rollout failed. Moved it into RendererConfig.build().
There was a problem hiding this comment.
Does it work if you call register inside recipes, like rl_grpo_muse_glimmer_30b_search_r1?
Also @felipemello1 has renderers accept non-HF-based tokenizers, so that we can remove transformers dependency?
There was a problem hiding this comment.
The PR has been up and approved since June, but never merged. I pinged them again about it. I will update when i hear back.
There was a problem hiding this comment.
@tianyu-l That's what it was doing before. register() was called inside rl_grpo_muse_glimmer_30b_search_r1(). It works for the controller, but the RolloutWorker only receives the serialized config and never imports the recipe module, so registration never happens in that process. That's the failure I hit after rebasing.
There was a problem hiding this comment.
I see. I'm OK with this being temporary solution. Please put a TODO and remove it later.
There was a problem hiding this comment.
They merged renderers not needing transformer anymore PrimeIntellect-ai/renderers#70
|
@HosseinKaviani-H , instead of search-R1, could we use dapo math? the search r1 relies on the embedding service, which has scalability issues. edit: i see that the renderer has tool call. So yeah, indeed search-r1 is a better case. If its not a big lift, could you also have a config for dapo math tested? |
@felipemello1 Yeah search-r1 was the only agentic task available in Titan and Glimmer is agentic. What would be the signal that you're interested in getting from dapo math? It's single turn no tools so it should be an easier one for the renderer, but Glimmer was built for agentic work so dapo might score low here. I can look into it more. |
With DAPO, you can run much longer contexts, e.g. 32k, and you can run multiple nodes. With search r1, the answer is like <200 tokens and if you try >1 node, the embedding server becomes a big bottleneck. |
a1bb096 to
61db2b2
Compare
61db2b2 to
3869b16
Compare
RL converts chat messages to tokens before generation and parses the completion
back into (content, reasoning_content, tool_calls) afterwards -- the rollout loop
needs that split to tell "call a tool and continue" apart from "final answer,
score it". Muse Glimmer has no renderer today, so it falls back to "auto", which
cannot parse its tool calls. This adds one.
Muse Glimmer uses the harmony chat format: an assistant turn is a sequence of
to=<recipient><|message|><body> channels, where to=self is private reasoning and
other recipients carry user-visible content. Tool calls are ATEM XML inside those
channels.
- models/muse_glimmer/renderer.py: MuseGlimmerRenderer, implementing the
renderers.Renderer Protocol -- render_ids, parse_response (harmony channel
splitting + ATEM extraction), get_stop_token_ids. register() installs it into
the renderers library's public registry (RENDERER_REGISTRY / _CONFIG_BY_NAME).
- models/muse_glimmer/atem.py: ATEM tool-call parse/render.
- experiments/rl/renderer.py: map "muse_glimmer" in _RENDERER_BY_MODEL.
Note on placement: every other model here resolves to a renderer in
PrimeIntellect-ai/renderers, and this is the first renderer to ship inside
torchtitan. It does so because Muse Glimmer is not in that library yet, and it
uses the library's supported extension path rather than forking it. Once the
renderer is upstreamed, register() can be deleted and only the
_RENDERER_BY_MODEL entry kept. Happy to invert that order and upstream first if
maintainers prefer.
Two behaviours are config knobs rather than hardcoded, both defaulting to the
conservative choice:
- reasoning_strength (default None -> template's own default). At the default
strength the model can spend a tight token budget entirely on reasoning and be
truncated before emitting a tool call or answer. Agentic tasks may want "low".
- answer_from_reasoning_fallback (default False). When the model emits only
reasoning -- no user-facing channel, no tool call -- optionally treat the last
reasoning line as the answer. Off by default because it promotes private
reasoning to user-visible content; useful for outcome-scored RL where an empty
content is unscoreable.
get_stop_token_ids deliberately excludes <|eom|> (200007): it ends the reasoning
channel, not the turn, so stopping there truncates before the tool call/answer.
Test Plan:
Resolution through the path RL actually uses:
register()
RendererConfig(name="muse_glimmer").build(tokenizer_path=<ckpt>)
-> MuseGlimmerRenderer, stop tokens [200008, 200001]
Round-trip against the released tokenizer:
- render_ids(messages, tools, add_generation_prompt=True) -> 391 tokens
- parse_response on a reasoning + ATEM tool-call completion recovers
reasoning_content "Need to look this up." and
tool_calls [("search", {"query": "Avatar director"})]
Config knobs:
- reasoning_strength=None vs "low" produce different prompts
- answer_from_reasoning_fallback False -> content ""; True -> content is the last
reasoning line
No regressions to other models: after register(), qwen3, gpt-oss, deepseek-v3 and
default (llama3) still resolve to their own renderers, and register() is
idempotent.
python -m torchtitan.models.muse_glimmer.atem -> ATEM parse/render round-trip
checks pass.
Known limitation: bridge_to_next_turn() returns None, the documented safe
fallback that makes the caller re-render the full prefix each turn. Correct but
leaves the multi-turn extension optimisation unimplemented; the byte-exact
version needs validation against real multi-tool generations.
Two fixes to render(), prompted by comparing against GptOssRenderer. 1. sampled_mask / is_content marked an assistant message's ENTIRE token span, header included. A harmony block is `<|start|> role [to=recipient] <|message|> body <|eom|>/<|eot|>`, and everything through <|message|> is template scaffolding the model never sampled -- so the trainer was putting policy gradient on scaffolding tokens. Now the header is excluded, matching how GptOssRenderer splits header / body / terminator. Note one assistant message can expand into several harmony blocks (a to=self reasoning channel plus one per tool call), so header state is tracked across the whole span rather than split once at the first <|message|>. On a 4-message reason -> tool-call -> tool-result -> answer exchange this drops 16 of 63 previously-sampled tokens. That fraction shrinks as generations get longer (headers are fixed-size, bodies are not), but it is systematic. 2. New retain_reasoning_in_history knob (default True, preserving current behaviour). The chat template emits a to=self channel for any assistant message carrying reasoning_content, so multi-turn history accumulates every turn's reasoning. Harmony-style models are often trained with prior analysis dropped from context -- gpt-oss does this via auto_drop_analysis -- so this makes the choice explicit rather than implicit. Left defaulting to the template's own behaviour because the right setting for Muse Glimmer needs checking against what it was trained on. Test Plan: Decoding the sampled span of a reason -> tool-call -> tool-result -> answer exchange now yields only model-authored text: 'Need to search.<|eom|><atem:function_calls>...</atem:function_calls><|eot|>James Cameron<|eot|>' with no <|start|>/<|message|> tokens present (asserted), where previously the assistant headers were included. retain_reasoning_in_history=True keeps the prior turn's reasoning in the re-render (142 tokens); False drops it (132), verified by substring. Unchanged: RendererConfig(name="muse_glimmer") resolves to MuseGlimmerRenderer, render_ids and parse_response round-trip reasoning + ATEM tool calls, and `python -m torchtitan.models.muse_glimmer.atem` passes.
Completes the path to running RL on Muse Glimmer. The model, parallelize and
state-dict adapter are already upstream; the renderer is the parent commit in
this stack. This adds the entry point that ties them together, so the model can
be launched the same way the Qwen3 and GPT-OSS Search-R1 recipes are:
python -m torchtitan.experiments.rl.train \
--module search_r1 --config rl_grpo_muse_glimmer_30b_search_r1
8 GPUs: 6 trainer (FSDP=3 x TP=2) + 2 generator (TP=2). Two settings are
model-specific and worth knowing before changing them:
* Generator TP <= 2. Muse Glimmer has 2 KV heads, so attention cannot be
tensor-split further; scale the trainer with FSDP rather than TP.
* ac_config=FullAC. Adam's m/v are allocated on the *first* optimizer.step(), so
per-GPU memory jumps ~8 bytes/param between step 1 and step 2 (~37 GB/GPU here,
sharded 6 ways). The default SelectiveAC OOMs at step 2; FullAC frees the
activation headroom that jump needs. This is the one setting most likely to be
"optimised" away by someone who has not hit the OOM.
varlen is used for both roles so trainer and generator share one ModelSpec.
Test Plan:
Ran the full loop on 8x H100 with this config (the released 30B checkpoint, the
stock RewardExactMatch rubric -- no leniency -- and a retrieval server):
Step 1 reward=0.62 entropy=0.27 grad_norm=3.46
Step 2 reward=0.28 entropy=0.32 grad_norm=2.81
Step 3 reward=0.66 entropy=0.26 grad_norm=3.87
Step 4 reward=0.42 entropy=0.30 grad_norm=4.21
0 OOM, 0 aborts, ~714 tok/s across ~148 concurrent requests. Reward carries real
variance rather than saturating, so every step produces a non-zero gradient.
Config assembly is exercised through the ConfigManager import path, and the
existing rl_grpo_qwen3_1_7b_search_r1 config still builds unchanged.
Separately verified on this stack: bitwise parity between the trainer and the
vLLM generator (3/3 subtests, max_delta = 0.00e+00), and that the state-dict
adapter reproduces the reference Q/K layout from the released HF export to ~9e-4.
Note: hf_assets_path follows the convention of the sibling configs and points at
torchtitan/experiments/rl/example_checkpoint/MuseGlimmer-30B; point it at your
own checkout of the released weights.
Three pre-commit failures from CI: - insert-license: renderer.py and atem.py were missing the standard header. - ufmt: reformatted renderer.py and config_registry.py. - pyrefly: `Could not find import of renderers.base`. pyrefly skips torchtitan/experiments (project-excludes), which is why the existing experiments/rl/renderer.py can import `renderers` freely -- but models/muse_glimmer/ is type-checked, and `renderers` is not installed in CI. Added `renderers.*` to replace-imports-with-any alongside the other optional dependencies (torchao, torchvision, torchcomms, ...), which is what it is: an RL-only optional dep. Verified after the changes: both the muse_glimmer and qwen3 Search-R1 configs still build, `python -m torchtitan.models.muse_glimmer.atem` passes, and `ufmt check` is clean on all four files.
…etup renderers validates in BaseRendererConfig.__pydantic_init_subclass__ that every non-base renderer-config field is declared as either a chat-template kwarg or a renderer-internal knob. MuseGlimmerRendererConfig declared neither set, so importing it raised TypeError on current renderers (the RL README installs it from @main, so there is no pinned version to have caught this). Declare both frozensets. reasoning_strength is a template field -- it is forwarded to apply_chat_template verbatim -- and the other two are internal knobs with no Jinja analogue. Older renderers ignore the ClassVars, so this works on both. Also document the NCCL_CTRAN_BACKENDS / NCCL_IB_DISABLE settings a single-node host without InfiniBand needs, and why they must be exported before launch rather than set inside the training script.
The renderer rendered prompts via tokenizer.apply_chat_template, which assumes the checkpoint ships a chat template. Not every Muse Glimmer checkpoint does: some carry it as a separate chat_template.jinja, others have none, and on those every rollout failed inside the first render with a generic per-rollout ERROR. Add a chat_template config field (path to a .jinja file, or inline source) and resolve it once at construction. When the tokenizer has no template and none is configured, raise there with a message naming the tokenizer and the fix, so the failure surfaces before training starts instead of once per rollout. Verified byte-identical token ids across all three paths (tokenizer-supplied, explicit path, inline source) on a checkpoint that ships a template.
The renderer lived in torchtitan/models/muse_glimmer, which pyrefly type-checks, and it imports renderers -- an optional RL-only dependency installed from git and absent in CI. That forced adding renderers.* to the project-level replace-imports-with-any list. All the other code importing renderers is under torchtitan/experiments, which project-excludes already covers, so this was the first import outside that boundary. RL is the renderer's only consumer, so move it to experiments/rl/models/muse_glimmer instead and revert pyproject.toml to upstream. No project-level setting changes, and the core model package stays importable without renderers installed.
Per review: the NCCL_CTRAN_BACKENDS / NCCL_IB_DISABLE settings are a property of a launching host that has no InfiniBand, not of TorchTitan RL, so they do not belong in the shared RL README.
The renderer called tokenizer.apply_chat_template, which made it depend on the checkpoint shipping a chat template and forced the loss mask to be recovered by diffing token prefixes across growing message lists. Every model-specific renderer in the renderers library builds tokens in Python instead -- only DefaultRenderer wraps Jinja -- including gpt_oss, which renders this same harmony format. Build the format directly. Rendering emits labelled spans, so loss attribution is exact by construction rather than inferred, and bridge_to_next_turn can now extend a sampled completion instead of returning None and forcing a re-render (which would re-tokenize what the generator produced). Add a parity test asserting the native render is byte-identical to apply_chat_template across 14 message shapes x generation-prompt x 3 configs, plus mask, bridge and tokenizer-compatibility checks: 117 cases. Verified the new path returns identical token ids and an identical loss mask to the Jinja implementation it replaces. Drops the chat_template config field, which existed only to feed apply_chat_template, and adds knowledge_cutoff / current_date -- the template substitutes today's date, so pinning it is what makes a run reproducible across days.
download_hf_assets.py writes to <local_dir>/<repo name>, so the canonical fetch for meta-models/Muse-Glimmer-30B lands in example_checkpoint/Muse-Glimmer-30B. The config asked for example_checkpoint/MuseGlimmer-30B, which the documented command never creates. Document the fetch in the example README the way dapo_math does.
TokenEnv.step reads bridged.token_ids, so returning a list made every multi-turn rollout fail with AttributeError the moment the first tool call came back. The unit test asserted the same wrong contract, so it passed; only an end-to-end run caught it. Match the interface the library defines and gpt_oss implements: return RenderedTokens with sampled_mask all False (a bridge produces a prompt), message_indices -1 over the carried prefix and the trailing generation prompt, and is_content marking message bodies. Use the library's own guards -- decline when there is no prior prompt, nothing to append, or an assistant turn in the extension, whose terminator depends on the role of the message after it. Use trim_to_turn_close so a completion truncated at max_tokens still gets a synthesized <|eot|>. Track is_body / is_generation_prompt on the emitted spans so the bridge can label tokens without re-deriving structure. Extend the tests to the real contract: array lengths, nothing sampled, the decline cases, and the truncated-completion path.
Per review: the plan to upstream the renderer and drop the custom registration was described in prose but not greppable. Add a TODO at the _RENDERER_BY_MODEL entry and another in the renderer's module docstring, naming the destination files so whoever picks it up does not have to rediscover them. Comments only.
Two divergences from how the renderers library's own renderers behave. is_content mirrored sampled_mask, so a user or tool message's body was reported as non-content. Both gpt_oss and qwen3 mark caller-supplied bodies as content on every role and exclude only header scaffolding; the terminator is content only on assistant turns, where the model emits its own stop token. Track it off the existing is_body span flag instead. Token ids are unchanged -- this only affects the is_content array. BaseRendererConfig.thinking_retention is the one field every renderer is expected to resolve and honour in its bridge, and this renderer ignored it: a caller could set it and nothing would happen. Resolve it in __init__ and decline to bridge when the policy requires a re-render. The implied default is "all", which is template-faithful -- the published chat template renders reasoning_content for every assistant turn with no query-boundary drop, unlike gpt-oss's auto_drop_analysis or qwen3's think-block stripping. Adds two tests; 121 pass, including the 84 byte-exactness cases, so neither change alters the token stream.
Upstream now builds the renderer inside a Monarch-spawned RolloutWorker (actors/rollout_worker.py -> rollout/rollouter.py setup_async). That process never imports the example's config module, so calling register() from the config registry left the worker with an unregistered renderer and every rollout failed with "No renderer config registered for name='muse_glimmer'". Register lazily in RendererConfig.build(), which is the code path that runs in the worker, and drop the now-redundant call from the config registry so there is a single registration site. Tagged with the same TODO as the _RENDERER_BY_MODEL entry: both go away once the renderer is upstreamed. The unit tests construct the renderer directly rather than through the actor, so they could not catch this; found by an end-to-end run after rebasing onto 14366c5.
Per review: the block explained why it exists and pointed at the TODO on the _RENDERER_BY_MODEL entry, but was not itself greppable as a TODO. Lead with it and state the removal condition -- the library registers its own renderers in _populate_registry(), so this hook disappears once Muse Glimmer is upstreamed. Comments only.
Upstream moved sequence-length settings out of AsyncLoopConfig.batcher and into TrainingConfig (num_tokens_per_microbatch_per_dp_rank / max_context_length), and dropped the Batcher import from this module. The Muse Glimmer config still passed batcher=Batcher.Config(...), so it failed to build with NameError: name 'Batcher' is not defined. Match what the Qwen3 configs in this file now do. Verified with a 100-step run on the public checkpoint: held-out 40 -> 92 of 200, in line with the three prior runs.
3869b16 to
54c2287
Compare
6 participants
Summary
Makes Muse Glimmer runnable in TorchTitan RL. The model, parallelize, and state-dict adapter are already upstream; this adds the two missing pieces : a renderer and a Search-R1 config. Therfore, it launches like the Qwen3 and GPT-OSS recipes:
python -m torchtitan.experiments.rl.train \ --module search_r1 --config rl_grpo_muse_glimmer_30b_search_r1Everything lives under
torchtitan/experiments/rl/. No core files, nopyproject.toml.What's here
Renderer (
experiments/rl/models/muse_glimmer/renderer.py,atem.py)RL needs messages -> tokens before generation, and tokens ->
(content, reasoning_content, tool_calls)after, so the rollout loop can tell "call a tool and continue" from "final answer, score it". Muse Glimmer had no renderer and fell back toauto, which can't parse its ATEM tool calls.Built natively in Python rather than wrapping
apply_chat_template, matching how every model-specific renderer in therendererslibrary works (onlyDefaultRendererwraps Jinja). This means it works on any checkpoint regardless of what the tokenizer ships, loss attribution is exact rather than recovered by diffing token prefixes, andbridge_to_next_turncan extend a sampled completion without re-rendering it.Tests (
experiments/rl/tests/test_muse_glimmer_renderer.py) - 119 cases, 84 of them asserting the native render is byte-identical toapply_chat_templateacross roles,tool shapes, reasoning states and config knobs. Treat it as the spec.Search-R1 config : 8 GPUs: 6 trainer (FSDP=3 x TP=2) + 2 generator (TP=2).
Two model-specific settings worth not "optimising" away:
ac_config=FullAC: Adam's m/v allocate on the firstoptimizer.step(), so memory jumps ~37 GB/GPU between step 1 and 2. SelectiveAC OOMs at step 2.Results
100 steps on 8xH100, released public checkpoint, HotpotQA, real dense retrieval over wiki-18 (e5 + FAISS HNSW). Stock
RewardExactMatch. Two independent runs. 0 OOM, 0 aborts.Matching the held-out questions one-to-one between step 0 and step 100 shows where the gain comes from and it reproduces across both runs:
So the gain is entirely questions it previously never completed. On the ~50 it could already answer it loses 7 which is a precision-for-coverage trade, reproducible across seeds. Noise floor on a 200-question pass is ~3 questions.
Also verified: bitwise parity trainer == vLLM generator (3/3 subtests,
max_delta = 0.00e+00), andrl_grpo_qwen3_1_7b_search_r1still builds unchanged.Notes
Retrieval is the ceiling, not the model. The answer appears in the retrieved passages 45% of the time at
topk=3and 59% attopk=10; held-out sits at 0.445. Two settings had to move together to get there.topk3 -> 10 and the rollout token budget. Neither worked alone.On entropy: ran the same config twice and entropy went up in one run and down in the other, so the direction isn't reproducible. It also tracks response length more than anything else. I don't have a clean explanation and wouldn't read much into the curve:
group_zero_std_fracandgrad_normwere the metrics that actually separated healthyruns from starved ones.
First in-repo renderer. Every other model resolves to one in PrimeIntellect-ai/renderers; this ships here because Muse Glimmer isn't in that library yet, using its supported extension path (
RENDERER_REGISTRY) rather than a fork. Once upstreamed,register()and both files can be deleted, leaving one_RENDERER_BY_MODELline.