[rl] Add Muse Glimmer renderer and Search-R1 RL config

torchtitan/experiments/rl/examples/search_r1/README.md

@@ -53,6 +53,19 @@ python <search-r1>/local_dense_retriever/retrieval_server.py \
 
 Override `message_env.search_url` / `message_env.topk` in the config if needed.
 
+### 3. Checkpoint
+Download the base checkpoint the config expects. `download_hf_assets.py` writes to a
+subdirectory named after the repo, which is the path in `hf_assets_path`:
+
+```bash
+python scripts/download_hf_assets.py \
+  --repo_id meta-models/Muse-Glimmer-30B \
+  --local_dir torchtitan/experiments/rl/example_checkpoint \
+  --all
+```
+
+Swap `--repo_id` for the model your config selects (e.g. `Qwen/Qwen3-1.7B`).
+
 ## Run
 
 ```bash

torchtitan/experiments/rl/examples/search_r1/config_registry.py

@@ -27,6 +27,7 @@
     ParallelismConfig,
     TrainingConfig,
 )
+from torchtitan.distributed.activation_checkpoint import FullAC
 from torchtitan.experiments.rl.actors.generator import (
     SamplingConfig,
     VLLMCudagraphConfig,
@@ -47,6 +48,10 @@
 from torchtitan.experiments.rl.observability.metrics import MetricsProcessor
 from torchtitan.experiments.rl.renderer import RendererConfig
 from torchtitan.experiments.rl.rollout.advantage import AdvantageEstimator
+from torchtitan.models.muse_glimmer import model_registry as muse_glimmer_model_registry
+from torchtitan.models.muse_glimmer.state_dict_adapter import (
+    MuseGlimmerStateDictAdapter,
+)
 from torchtitan.models.qwen3 import model_registry
 
 
@@ -245,3 +250,96 @@ def rl_grpo_qwen3_30b_a3b_deepep_search_r1_perf() -> Controller.Config:
     # from this scheduler limit, CUDA graph capture sizes, CP, and SP.
     config.generator.max_num_batched_tokens = 2048  # TODO: TBD
     return config
+
+
+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
+    server on spare capacity. Requires a running retrieval server and the QA parquet
+    data; see ``README.md``.
+
+    Two constraints are specific to this model:
+
+    * **Generator TP <= 2.** Muse Glimmer has 2 KV heads, so attention cannot be
+      tensor-split further. Scale the trainer with FSDP rather than TP.
+    * **Full activation checkpointing is required.** Adam's m/v are allocated on the
+      *first* ``optimizer.step()``, so per-GPU memory jumps by roughly 8 bytes/param
+      between step 1 and step 2 (~37 GB/GPU here, sharded 6 ways). With the default
+      ``SelectiveAC`` that jump OOMs at step 2; ``FullAC`` frees the activation
+      headroom it needs.
+
+    varlen attention is used for both roles so the trainer and the vLLM generator run
+    one ModelSpec. The state-dict adapter handles the HF checkpoint's Q/K RoPE layout
+    on load, and the renderer (registered below) handles Muse Glimmer's harmony chat
+    format and ATEM tool calls.
+    """
+    # Muse Glimmer's renderer ships in torchtitan rather than the `renderers` library;
+    # registering makes RendererConfig(name="muse_glimmer") resolve it.
+
+    model_spec = muse_glimmer_model_registry("30B", attn_backend="varlen")
+    model_spec = dataclasses.replace(
+        model_spec, state_dict_adapter=MuseGlimmerStateDictAdapter
+    )
+
+    return Controller.Config(
+        model_spec=model_spec,
+        hf_assets_path="torchtitan/experiments/rl/example_checkpoint/Muse-Glimmer-30B",
+        async_loop=AsyncLoopConfig(
+            num_training_steps=500,
+            num_prompts_per_train_step=8,
+            num_samples_per_prompt=8,
+            validation=ValidationConfig(num_samples=500),
+        ),
+        compile=CompileConfig(enable=False),
+        rollouter=SearchR1Rollouter.Config(
+            worker=SearchR1Worker.Config(
+                advantage=AdvantageEstimator.Config(should_std_normalize=True),
+            ),
+        ),
+        renderer=RendererConfig(name="muse_glimmer", enable_thinking=True),
+        metrics=MetricsProcessor.Config(enable_wandb=True),
+        trainer=PolicyTrainer.Config(
+            optimizer=default_adamw(lr=1e-6),
+            lr_scheduler=LRSchedulersContainer.Config(
+                warmup_steps=2, decay_type="linear", min_lr_factor=1.0
+            ),
+            training=TrainingConfig(
+                num_tokens_per_microbatch_per_dp_rank=4096,
+                max_context_length=4096,
+            ),
+            ac_config=FullAC.Config(),
+            parallelism=ParallelismConfig(
+                data_parallel_shard_degree=3,
+                tensor_parallel_degree=2,
+            ),
+            checkpoint=CheckpointManager.Config(
+                enable=True,
+                initial_load_in_hf=True,  # first run loads HF; restarts resume from DCP
+                interval=50,
+                last_save_model_only=False,
+                keep_latest_k=3,
+            ),
+            loss=ChunkedLossWrapper.Config(
+                num_chunks=8,
+                loss_fn=DAPOLoss.Config(
+                    ratio_clip_low=0.2,
+                    ratio_clip_high=0.28,
+                ),
+            ),
+        ),
+        generator=VLLMGenerator.Config(
+            model_dtype="bfloat16",
+            parallelism=InferenceParallelismConfig(
+                data_parallel_degree=1,
+                tensor_parallel_degree=2,  # <= 2 KV heads
+            ),
+            cudagraph=VLLMCudagraphConfig(enable=False),
+            checkpoint=CheckpointManager.Config(enable=False),
+            sampling=SamplingConfig(
+                temperature=1.0,
+                top_p=1.0,
+                max_tokens=4096,
+            ),
+        ),
+    )

torchtitan/experiments/rl/models/muse_glimmer/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.

torchtitan/experiments/rl/models/muse_glimmer/atem.py

@@ -0,0 +1,90 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Muse Glimmer ATEM tool-call parse/render (the novel core of the Muse Glimmer renderer).
+
+Muse Glimmer emits tool calls in Anthropic-style ATEM XML inside the harmony envelope:
+
+    <atem:function_calls>
+    <atem:invoke name="search">
+    <atem:parameter name="query">who wrote Blade Runner</atem:parameter>
+    </atem:invoke>
+    </atem:function_calls>
+
+parse: model text -> [{"name", "arguments"}]  (what env.step() needs)
+render: a tool call -> ATEM text  (what goes back into the prompt)
+"""
+
+from __future__ import annotations
+
+import json
+import re
+
+_FUNCTION_CALLS = re.compile(
+    r"<atem:function_calls>(.*?)</atem:function_calls>", re.DOTALL
+)
+_INVOKE = re.compile(
+    r'<atem:invoke name="(?P<name>[^"]+)">(?P<body>.*?)</atem:invoke>', re.DOTALL
+)
+_PARAMETER = re.compile(
+    r'<atem:parameter name="(?P<key>[^"]+)">(?P<value>.*?)</atem:parameter>', re.DOTALL
+)
+
+
+def parse_atem_tool_calls(text: str) -> list[dict]:
+    """Parse every ATEM tool call in `text` into `[{"name", "arguments"}]`.
+
+    Values are JSON-decoded when possible (dicts/lists/numbers/bools), else kept
+    as the raw string. Supports multiple parallel invokes in one block.
+    """
+    calls: list[dict] = []
+    for block in _FUNCTION_CALLS.findall(text):
+        for invoke in _INVOKE.finditer(block):
+            arguments: dict = {}
+            for param in _PARAMETER.finditer(invoke.group("body")):
+                raw = param.group("value")
+                try:
+                    arguments[param.group("key")] = json.loads(raw)
+                except (json.JSONDecodeError, ValueError):
+                    arguments[param.group("key")] = raw
+            calls.append({"name": invoke.group("name"), "arguments": arguments})
+    return calls
+
+
+def render_atem_tool_call(name: str, arguments: dict) -> str:
+    """Render one tool call as ATEM text (matches Muse Glimmer's chat template)."""
+    lines = ["<atem:function_calls>", f'<atem:invoke name="{name}">']
+    for key, value in arguments.items():
+        if isinstance(value, bool):
+            sval = "true" if value else "false"
+        elif value is None:
+            sval = "null"
+        elif isinstance(value, (dict, list)):
+            sval = json.dumps(value)
+        else:
+            sval = str(value)
+        lines.append(f'<atem:parameter name="{key}">{sval}</atem:parameter>')
+    lines += ["</atem:invoke>", "</atem:function_calls>"]
+    return "\n".join(lines)
+
+
+if __name__ == "__main__":
+    # round-trip self-test (no deps, no GPU)
+    sample = (
+        'thinking...\n<atem:function_calls>\n<atem:invoke name="search">\n'
+        '<atem:parameter name="query">who wrote Blade Runner</atem:parameter>\n'
+        "</atem:invoke>\n</atem:function_calls>"
+    )
+    calls = parse_atem_tool_calls(sample)
+    assert calls == [
+        {"name": "search", "arguments": {"query": "who wrote Blade Runner"}}
+    ], calls
+    # no tool call -> empty (this is how env.step() detects "final answer")
+    assert parse_atem_tool_calls("Philip K. Dick") == []
+    # render -> parse round-trip
+    rendered = render_atem_tool_call("search", {"query": "x", "topk": 3})
+    assert parse_atem_tool_calls(rendered)[0]["arguments"]["topk"] == 3
+    print("muse_glimmer atem: all checks passed")