[rl] Pin renderers 0.1.11, render with TorchTitan's tokenizer, take the renderer's typed config directly

.github/CODEOWNERS

@@ -9,6 +9,5 @@
 # Exclude the experiments directory by adding a pattern without owners
 /torchtitan/experiments/
 
-/torchtitan/experiments/forge/ @felipemello1 @tianyu-l @wwwjn @fegin
 
 /torchtitan/experiments/graph_trainer/ @SherlockNoMad @xmfan @aditvenk @sanketpurandare @IvanKobzarev @tianyu-l

.github/workflows/integration_test_b200.yaml

@@ -2,15 +2,10 @@ name: B200 Integration
 
 on:
   push:
-    branches: [ main ]
     tags:
       - ciflow/b200/*
-    paths:
-      - 'torchtitan/models/kimi_k3/**'
-      - 'torchtitan_recipes/tests/b200.py'
-      - 'tests/integration_tests/b200.py'
-      - '.github/workflows/integration_test_b200.yaml'
-  workflow_dispatch:
+  schedule:
+    - cron: '0 0 * * *'
 
 concurrency:
   group: unit-test-${{ github.workflow }}-${{ github.ref == 'refs/heads/main' && github.run_number || github.ref }}

.github/workflows/integration_test_h100.yaml

@@ -4,6 +4,8 @@ on:
   push:
     tags:
       - ciflow/h100.8/*
+  schedule:
+    - cron: '0 0 * * *'
 
 concurrency:
   group: unit-test-${{ github.workflow }}-${{ github.ref == 'refs/heads/main' && github.run_number || github.ref }}

README.md

@@ -39,7 +39,7 @@ The Guiding Principles when building `torchtitan`
 * Minimal changes to the model code when applying multi-dimensional parallelism.
 * Bias towards a clean, minimal codebase while providing basic reusable / swappable components.
 
-`torchtitan` has been showcasing PyTorch's latest distributed training features, via support for pretraining Llama 3.1 LLMs of various sizes.
+`torchtitan` showcases PyTorch's latest distributed training features across multiple model families. Core models include Llama 3, Qwen3 / 3.5 / 3.8, DeepSeek V3 / V4, GPT-OSS, Kimi K2.7 / K3, Muse Glimmer, and Flux.
 
 ## Contributing
 
@@ -53,7 +53,7 @@ We look forward to your contributions!
 | Hardware | Integration Tests | Unit Tests |
 | --- | --- | --- |
 | CPU | - | [![CPU Unit Test](https://github.com/pytorch/torchtitan/actions/workflows/unit_test_cpu.yaml/badge.svg?branch=main)](https://github.com/pytorch/torchtitan/actions/workflows/unit_test_cpu.yaml?query=branch%3Amain) |
-| NVIDIA GPU | [![Integration Tests](https://github.com/pytorch/torchtitan/actions/workflows/integration_test.yaml/badge.svg?branch=main)](https://github.com/pytorch/torchtitan/actions/workflows/integration_test.yaml?query=branch%3Amain) [![H100 Tests](https://github.com/pytorch/torchtitan/actions/workflows/integration_test_h100.yaml/badge.svg?branch=main)](https://github.com/pytorch/torchtitan/actions/workflows/integration_test_h100.yaml?query=branch%3Amain) | [![GPU Unit Tests](https://github.com/pytorch/torchtitan/actions/workflows/unit_test_gpu.yaml/badge.svg?branch=main)](https://github.com/pytorch/torchtitan/actions/workflows/unit_test_gpu.yaml?query=branch%3Amain) |
+| NVIDIA GPU | [![Integration Tests](https://github.com/pytorch/torchtitan/actions/workflows/integration_test.yaml/badge.svg?branch=main)](https://github.com/pytorch/torchtitan/actions/workflows/integration_test.yaml?query=branch%3Amain) [![H100 Tests](https://github.com/pytorch/torchtitan/actions/workflows/integration_test_h100.yaml/badge.svg?branch=main)](https://github.com/pytorch/torchtitan/actions/workflows/integration_test_h100.yaml?query=branch%3Amain) [![B200 Tests](https://github.com/pytorch/torchtitan/actions/workflows/integration_test_b200.yaml/badge.svg?branch=main)](https://github.com/pytorch/torchtitan/actions/workflows/integration_test_b200.yaml?query=branch%3Amain) | [![GPU Unit Tests](https://github.com/pytorch/torchtitan/actions/workflows/unit_test_gpu.yaml/badge.svg?branch=main)](https://github.com/pytorch/torchtitan/actions/workflows/unit_test_gpu.yaml?query=branch%3Amain) |
 
 ## Llama 3.1 training
 
@@ -70,7 +70,7 @@ We look forward to your contributions!
    - [Interoperable checkpoints](docs/checkpoint.md) which can be loaded directly into [`torchtune`](https://github.com/pytorch/torchtune) for fine-tuning
 5. `torch.compile` support
 6. [Float8](https://discuss.pytorch.org/t/distributed-w-torchtitan-enabling-float8-all-gather-in-fsdp2/209323) support ([how-to](torchtitan/components/quantization/float8.md))
-7. [MXFP8 training for dense and MoE models](torchtitan/components/quantization/mxfp8.md) on Blackwell GPUs.
+7. [MXFP8 training for dense and MoE models](torchtitan/components/quantization/mxfp8/README.md) on Blackwell GPUs.
 8. Supervised Fine-Tuning (SFT) with chat-formatted datasets
 9. DDP and HSDP
 10. [TorchFT](https://github.com/pytorch/torchft) integration
@@ -84,10 +84,10 @@ We look forward to your contributions!
 17. All options easily configured in [Python](torchtitan/config/README.md) with `--module` and `--config` CLI flags
 18. Structured logging: per-rank trace of key training phases; (see [`torchtitan/observability/structured_logger/README.md`](torchtitan/observability/structured_logger/README.md))
 19. [Helper scripts](scripts/) to
-    - download tokenizers from Hugging Face
-    - convert original Llama 3 checkpoints into the expected DCP format
-    - estimate FSDP/HSDP memory usage without materializing the model
-    - run distributed inference with Tensor Parallel
+    - download tokenizers and other Hugging Face assets (`scripts/download_hf_assets.py`)
+    - convert checkpoints between Hugging Face and DCP formats (`scripts/checkpoint_conversion/`)
+    - compare training losses across commits or configs (`scripts/loss_compare.py`)
+    - run vLLM inference with TorchTitan models (`torchtitan/experiments/rl/generate.py`)
 
 We report [performance](benchmarks/llama3_h100_202412_torchtitan.md) on up to 512 GPUs, and verify [loss converging](docs/converging.md) correctness of various techniques.
 

tests/assets/losses/fake_pg/deepseek_v3_a10g.txt

@@ -2,13 +2,13 @@
 # ngpu: 8
 # parallelism: FSDP=8, EP=8
 # step loss grad_norm
-1 8.013261795043945 2.311385154724121
-2 7.456464767456055 2.8374972343444824
-3 7.099932670593262 3.907268524169922
-4 6.913008689880371 4.062598705291748
-5 6.833977699279785 3.8018505573272705
-6 6.724702835083008 3.8861584663391113
-7 6.630510330200195 3.739548683166504
-8 6.443962097167969 3.7964563369750977
-9 6.4159955978393555 3.8824899196624756
-10 6.394824028015137 3.693788766860962
+1 8.013107299804688 2.3114163875579834
+2 7.456279277801514 2.8375244140625
+3 7.100132942199707 3.907575845718384
+4 6.913045883178711 4.062716007232666
+5 6.834068298339844 3.8019232749938965
+6 6.724488258361816 3.886246681213379
+7 6.630580425262451 3.7395145893096924
+8 6.4435601234436035 3.796462059020996
+9 6.415536880493164 3.882359743118286
+10 6.3946332931518555 3.6937999725341797

tests/assets/losses/real_pg/deepseek_v3_cp_pp_a10g.txt

@@ -2,13 +2,13 @@
 # ngpu: 8
 # parallelism: FSDP=2, CP=2, PP=2, EP=4
 # step loss grad_norm
-1 8.122339248657227 3.625427722930908
-2 6.290388107299805 4.343588829040527
-3 4.845731735229492 3.05355167388916
-4 4.65423583984375 2.890531539916992
-5 4.4774627685546875 2.445112705230713
-6 4.265349388122559 2.1282083988189697
-7 4.171058654785156 2.0247530937194824
-8 4.110886573791504 1.9594563245773315
-9 4.0594587326049805 1.7155301570892334
-10 3.943784713745117 1.7355142831802368
+1 8.122316360473633 3.625668525695801
+2 6.290362358093262 4.343724250793457
+3 4.8458027839660645 3.0538036823272705
+4 4.654213905334473 2.8905367851257324
+5 4.477372169494629 2.4450645446777344
+6 4.265354156494141 2.1282594203948975
+7 4.171051979064941 2.024583339691162
+8 4.110989093780518 1.9595190286636353
+9 4.0594611167907715 1.7154442071914673
+10 3.943953275680542 1.7357909679412842

tests/integration_tests/b200.py

@@ -18,4 +18,18 @@ def build_b200_tests_list() -> list[OverrideDefinitions]:
             test_name="kimi_k3_mm_fsdp",
             ngpu=2,
         ),
+        # TODO: re-enable once the B200 job installs torchao. It currently
+        # installs only nightly torch/torchvision, requirements.txt and
+        # requirements-vlm.txt, none of which pull torchao in, so MXFP8Linear
+        # cannot import and MXFP8LinearConverter raises at construction.
+        # A plain `pip install torchao` is not enough either: the 32x32
+        # swizzled cast kernels landed in pytorch/ao#4777 and are unreleased as
+        # of v0.18.0, so this needs a source install or a later release.
+        OverrideDefinitions(
+            configs=[recipes.llama3_debugmodel_mxfp8_fsdp2],
+            test_descr="MXFP8 linear with an FSDP-managed weight cache",
+            test_name="mxfp8_linear_fsdp",
+            ngpu=2,
+            disabled=True,
+        ),
     ]

tests/unit_tests/cpu/test_checkpoint.py

@@ -23,6 +23,7 @@
 import torch.nn as nn
 from torch.distributed.checkpoint.state_dict_saver import AsyncSaveResponse
 from torch.utils.data import DataLoader
+
 from torchtitan.components.checkpointer.base import (
     BaseCheckpointManager,
     CheckpointStorage,
@@ -35,6 +36,7 @@
     AsyncMode,
     CheckpointManager,
 )
+from torchtitan.components.quantization._fsdp_tensor import _ShardedFSDPTensor
 from torchtitan.config import Function
 
 
@@ -739,6 +741,48 @@ def test_async_save_calls_maybe_wait_for_saving(
         new_future = manager.save_future
         new_future.result.assert_not_called()
 
+    @mock.patch("torchtitan.components.checkpointer.dcp.dist.new_group")
+    def test_purge_runs_before_this_step_save_is_issued(self, _mock_new_group):
+        trainer_config = DummyTrainerConfig(dump_folder=self.trainer_config.dump_folder)
+        checkpoint_config = trainer_config.checkpoint
+        checkpoint_config.async_mode = "async"
+        manager = CheckpointManager(
+            dataloader=self.data_loader,
+            model_parts=self.model_parts,
+            optimizers=self.optimizers,
+            lr_schedulers=self.lr_schedulers,
+            states=self.states,
+            config=checkpoint_config,
+            sd_adapter=None,
+            base_folder=self.trainer_config.dump_folder,
+        )
+        save_future: Future[None] = Future()
+        calls = []
+
+        with (
+            mock.patch.object(
+                manager,
+                "_purge_stale_checkpoints",
+                side_effect=lambda *, saving_step: calls.append(("purge", saving_step)),
+            ),
+            mock.patch.object(
+                manager,
+                "dcp_save",
+                side_effect=lambda *args, **kwargs: (
+                    calls.append("save"),
+                    save_future,
+                )[1],
+            ),
+            mock.patch(
+                "torchtitan.components.checkpointer.dcp.GarbageCollection.collect"
+            ),
+        ):
+            self.assertTrue(manager.save(curr_step=10))
+
+        self.assertEqual([("purge", 10), "save"], calls)
+        save_future.set_result(None)
+        manager.close()
+
     @mock.patch("torch.distributed.get_rank", return_value=0)
     @mock.patch.object(dist_checkpoint, "save")
     def test_enable_first_step_checkpoint(self, mock_save, mock_rank):
@@ -1247,7 +1291,7 @@ def setUp(self):
         self.addCleanup(shutil.rmtree, self.root, ignore_errors=True)
 
         self.manager = CheckpointManager.__new__(CheckpointManager)
-        self.manager.keep_latest_k = 1
+        self.manager.keep_latest_k = 2
         self.manager.folder = self.root
         self.manager._storage = _FilesystemCheckpointStorage()
         self.manager.purge_thread = mock.sentinel.purge_thread
@@ -1261,19 +1305,24 @@ def _write_checkpoint(self, name, *, complete=True):
                 pass
 
     @mock.patch("torch.distributed.get_rank", return_value=0)
-    def test_only_queues_complete_canonical_checkpoints(self, _rank):
+    def test_only_queues_stale_canonical_directories(self, _rank):
         self._write_checkpoint("step-1")
         self._write_checkpoint("step-2")
         self._write_checkpoint("step-100.backup")
         self._write_checkpoint("foo-step-101")
         self._write_checkpoint("step-0200")
         self._write_checkpoint("step-300", complete=False)
 
-        self.manager._purge_stale_checkpoints()
+        self.manager._purge_stale_checkpoints(saving_step=400)
 
-        self.manager.purge_queue.put.assert_called_once_with(
-            os.path.join(self.root, "step-1")
+        self.assertEqual(
+            [
+                mock.call(os.path.join(self.root, "step-1")),
+                mock.call(os.path.join(self.root, "step-300")),
+            ],
+            self.manager.purge_queue.put.call_args_list,
         )
+        self.assertTrue(os.path.exists(os.path.join(self.root, "step-300")))
 
 
 class TestSharedDiscoveryAndRetention(unittest.TestCase):
@@ -1305,14 +1354,46 @@ def test_bodies_are_defined_on_the_base(self):
                 self.assertIn(name, vars(BaseCheckpointManager))
 
     @mock.patch("torch.distributed.get_rank", return_value=0)
-    def test_purge_keeps_k_because_dcp_purges_after_saving(self, _rank):
-        # This manager purges once its checkpoint is already on disk, so it
-        # reserves nothing and keeps the full k.
+    def test_purge_reserves_a_slot_for_the_upcoming_save(self, _rank):
         manager = self._manager(keep_latest_k=2, entries=["step-1", "step-2", "step-3"])
 
-        manager._purge_stale_checkpoints()
+        manager._purge_stale_checkpoints(saving_step=4)
 
-        self.assertEqual({"/checkpoint/step-1"}, self._purged(manager))
+        self.assertEqual(
+            {"/checkpoint/step-1", "/checkpoint/step-2"},
+            self._purged(manager),
+        )
+
+    @mock.patch("torch.distributed.get_rank", return_value=0)
+    def test_incomplete_directory_cannot_evict_a_valid_checkpoint(self, _rank):
+        # An interrupted save leaves a step-N directory with no metadata. If it
+        # occupied a slot it would push a checkpoint we can actually resume from
+        # out of the retained set.
+        manager = self._manager(keep_latest_k=2, entries=["step-1", "step-2", "step-3"])
+        manager._storage.isfile.side_effect = lambda path: "step-3" not in path
+
+        manager._purge_stale_checkpoints(saving_step=4)
+
+        # k-1 valid ones stay (step-2), the incomplete step-3 is purged rather
+        # than counted, and step-1 falls out normally.
+        self.assertEqual(
+            {"/checkpoint/step-1", "/checkpoint/step-3"},
+            self._purged(manager),
+        )
+        manager._storage.remove.assert_not_called()
+
+    @mock.patch("torch.distributed.get_rank", return_value=0)
+    def test_abandoned_directories_are_queued_for_purge(self, _rank):
+        manager = self._manager(keep_latest_k=2, entries=["step-1", "step-2"])
+        manager._storage.isfile.return_value = False
+
+        manager._purge_stale_checkpoints(saving_step=3)
+
+        self.assertEqual(
+            {"/checkpoint/step-1", "/checkpoint/step-2"},
+            self._purged(manager),
+        )
+        manager._storage.remove.assert_not_called()
 
     @mock.patch("torch.distributed.get_rank", return_value=0)
     def test_purge_keeps_exempt_checkpoints_outside_latest_k(self, _rank):
@@ -1322,14 +1403,14 @@ def test_purge_keeps_exempt_checkpoints_outside_latest_k(self, _rank):
         )
         manager.purge_exempt = Function.Config(fn=lambda step: step % 2 == 0).build()
 
-        manager._purge_stale_checkpoints()
+        manager._purge_stale_checkpoints(saving_step=6)
 
         self.assertEqual(
             {"/checkpoint/step-1", "/checkpoint/step-3"},
             self._purged(manager),
         )
 
-    def test_parse_step_accepts_only_canonical_names(self):
+    def test_parse_step_accepts_only_canonical_published_names(self):
         manager = CheckpointManager.__new__(CheckpointManager)
 
         self.assertEqual(0, manager._parse_step("step-0"))
@@ -1338,6 +1419,16 @@ def test_parse_step_accepts_only_canonical_names(self):
             with self.subTest(name=name):
                 self.assertIsNone(manager._parse_step(name))
 
+    @mock.patch("torch.distributed.get_rank", return_value=0)
+    def test_purge_preserves_the_checkpoint_currently_being_saved(self, _rank):
+        manager = self._manager(keep_latest_k=2, entries=["step-1", "step-2"])
+        manager._storage.isfile.return_value = False
+
+        manager._purge_stale_checkpoints(saving_step=2)
+
+        self.assertEqual({"/checkpoint/step-1"}, self._purged(manager))
+        manager._storage.remove.assert_not_called()
+
     def test_valid_checkpoint_accepts_dcp_or_hf_markers(self):
         manager = CheckpointManager.__new__(CheckpointManager)
         manager._storage = mock.Mock(spec=CheckpointStorage)
@@ -1440,6 +1531,31 @@ def _split(module, state_dict, prefix, local_metadata):
         # ... and the in-place refresh picked up the updated parameter.
         self.assertTrue(torch.all(sd2["a"] == 1.0))
 
+    def test_fsdp_unsharded_tensor_checkpoint(self):
+        class FSDPWeight(_ShardedFSDPTensor):
+            pass
+
+        class WrappedBuffer(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("w", FSDPWeight(torch.zeros(4)))
+
+        model = WrappedBuffer()
+        wrapper = ModelWrapper(model)
+        state_dict = wrapper.state_dict()
+        cached_weight = state_dict["w"]
+
+        with torch.no_grad():
+            model.w._tensor.fill_(2.0)
+
+        refreshed = wrapper.state_dict()
+        self.assertIs(refreshed, state_dict)
+        self.assertIs(refreshed["w"], cached_weight)
+        self.assertTrue(torch.all(refreshed["w"]._tensor == 2.0))
+
+        wrapper.load_state_dict({"w": torch.full((4,), 3.0)})
+        self.assertTrue(torch.all(model.w._tensor == 3.0))
+
 
 if __name__ == "__main__":
     unittest.main()

tests/unit_tests/cpu/test_deepseek_v4_flops.py

@@ -0,0 +1,31 @@
+# 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.
+
+import unittest
+
+import torch
+
+from torchtitan.models.deepseek_v4 import model_registry
+
+
+class TestDeepSeekV4Flops(unittest.TestCase):
+    def test_flash_mtp_4k_model_flops(self):
+        model_config = model_registry(
+            "deepseek_v4_flash",
+            n_mtp_layers=1,
+        ).model
+
+        with torch.device("meta"):
+            model = model_config.build()
+
+        self.assertEqual(
+            model_config.get_nparams_and_flops(model, seq_len=4096),
+            (290_942_278_866, 92_762_352_876),
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()