diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 9ce458a529..f829cb8a6e 100644 --- a/.github/CODEOWNERS +++ b/.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 diff --git a/.github/workflows/integration_test_b200.yaml b/.github/workflows/integration_test_b200.yaml index 01548fce00..06877598f9 100644 --- a/.github/workflows/integration_test_b200.yaml +++ b/.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 }} diff --git a/.github/workflows/integration_test_h100.yaml b/.github/workflows/integration_test_h100.yaml index 0fbd0fa27d..6a88de4cae 100644 --- a/.github/workflows/integration_test_h100.yaml +++ b/.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 }} diff --git a/README.md b/README.md index 8a7503c78a..bdb04900b8 100644 --- a/README.md +++ b/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. diff --git a/assets/images/mxfp8_32x32_vs_bf16_loss.png b/assets/images/mxfp8_32x32_vs_bf16_loss.png new file mode 100644 index 0000000000..766a0a305e Binary files /dev/null and b/assets/images/mxfp8_32x32_vs_bf16_loss.png differ diff --git a/assets/images/mxfp8_with_loss.png b/assets/images/mxfp8_with_loss.png deleted file mode 100644 index 47e2967aed..0000000000 Binary files a/assets/images/mxfp8_with_loss.png and /dev/null differ diff --git a/tests/assets/losses/fake_pg/deepseek_v3_a10g.txt b/tests/assets/losses/fake_pg/deepseek_v3_a10g.txt index d84290ce14..45e6c22acc 100644 --- a/tests/assets/losses/fake_pg/deepseek_v3_a10g.txt +++ b/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 diff --git a/tests/assets/losses/real_pg/deepseek_v3_cp_pp_a10g.txt b/tests/assets/losses/real_pg/deepseek_v3_cp_pp_a10g.txt index 4dae448070..b8368c7a4c 100644 --- a/tests/assets/losses/real_pg/deepseek_v3_cp_pp_a10g.txt +++ b/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 diff --git a/tests/integration_tests/b200.py b/tests/integration_tests/b200.py index c0ca6c9b19..4ba8cab4a6 100644 --- a/tests/integration_tests/b200.py +++ b/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, + ), ] diff --git a/tests/unit_tests/cpu/test_checkpoint.py b/tests/unit_tests/cpu/test_checkpoint.py index 39c4f534e8..8e7761b8fa 100644 --- a/tests/unit_tests/cpu/test_checkpoint.py +++ b/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,7 +1305,7 @@ 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") @@ -1269,11 +1313,16 @@ def test_only_queues_complete_canonical_checkpoints(self, _rank): 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() diff --git a/tests/unit_tests/cpu/test_deepseek_v4_flops.py b/tests/unit_tests/cpu/test_deepseek_v4_flops.py new file mode 100644 index 0000000000..4313433515 --- /dev/null +++ b/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() diff --git a/tests/unit_tests/cpu/test_integration_test_definitions.py b/tests/unit_tests/cpu/test_integration_test_definitions.py index a845e4db71..1e0686406c 100644 --- a/tests/unit_tests/cpu/test_integration_test_definitions.py +++ b/tests/unit_tests/cpu/test_integration_test_definitions.py @@ -92,7 +92,10 @@ def test_h100_tests_are_registered_in_separate_suite() -> None: def test_b200_tests_are_registered_in_separate_suite() -> None: - assert {test.test_name for test in build_b200_tests_list()} == {"kimi_k3_mm_fsdp"} + assert {test.test_name for test in build_b200_tests_list()} == { + "kimi_k3_mm_fsdp", + "mxfp8_linear_fsdp", + } assert "kimi_k3_mm_fsdp" not in { test.test_name for test in build_model_tests_list() } diff --git a/tests/unit_tests/cpu/test_quantization.py b/tests/unit_tests/cpu/test_quantization.py index c36d3ded30..40c5a614d6 100644 --- a/tests/unit_tests/cpu/test_quantization.py +++ b/tests/unit_tests/cpu/test_quantization.py @@ -19,10 +19,15 @@ from torchtitan.components.data.sources import HuggingFaceRandomAccessSource from torchtitan.components.quantization import Float8Linear from torchtitan.components.quantization.float8 import _get_float8_grouped_experts_cls -from torchtitan.components.quantization.mx import _get_mxfp8_grouped_experts_cls +from torchtitan.components.quantization.mxfp8.converter import ( + _get_mxfp8_grouped_experts_cls, + MXFP8Linear, + MXFP8LinearConverter, +) from torchtitan.components.quantization.utils import has_quantization from torchtitan.config import ConfigManager from torchtitan.models.common.decoder_sharding import colwise_config, rowwise_config +from torchtitan.models.common.feed_forward import FeedForward from torchtitan.models.common.linear import Linear from torchtitan.models.common.moe import GroupedExperts from torchtitan.models.gpt_oss.moe import GptOssGroupedExperts @@ -448,3 +453,226 @@ def test_float8_grouped_experts_dcp_round_trip_needs_no_safe_globals(tmp_path): source.parameters(), target.parameters(), strict=True ): torch.testing.assert_close(target_parameter, source_parameter) + + +def test_mxfp8_linear_validates_config_and_installs_weight_wrapper(): + pytest.importorskip("torchao") + if MXFP8Linear is None: + pytest.skip("torchao MXFP8Linear is unavailable") + from torchtitan.components.quantization._fsdp_tensor import _UnshardedFSDPTensor + from torchtitan.components.quantization.mxfp8.tensor import ( + _LinearShardedTensorWithMXFP8Compute, + ) + + with pytest.raises(ValueError, match="in_features divisible by 32"): + MXFP8Linear.Config(in_features=127, out_features=128) + with pytest.raises(ValueError, match="out_features divisible by 32"): + MXFP8Linear.Config(in_features=128, out_features=127) + with pytest.raises( + ValueError, + match="input_activation_format_for_backward must be one of", + ): + MXFP8Linear.Config( + in_features=128, + out_features=128, + input_activation_format_for_backward="missing", + ) + + for sharding_config in (colwise_config(), rowwise_config()): + linear = MXFP8Linear.Config( + in_features=128, + out_features=128, + bias=False, + sharding_config=sharding_config, + ).build() + assert linear._sharding_config is not None + # The wrapper is installed at construction, so no caller has to opt + # in. Until a data parallel implementation drives its lifecycle it is + # the sharded state, which holds the BF16 weight; the unsharded tensor + # is a separate type the post-all-gather hook produces. + assert isinstance(linear.weight, _LinearShardedTensorWithMXFP8Compute) + assert not isinstance(linear.weight, _UnshardedFSDPTensor) + + +def test_mxfp8_linear_rejects_the_partial_dtensor_backend(): + """MXFP8 needs the spmd_types backend to survive tensor parallelism. + + The matmul is an opaque autograd function, so DTensor has no sharding + strategy for it and propagation fails on the storage-free unsharded tensor. + spmd_types annotates the function instead. + """ + pytest.importorskip("torchao") + if MXFP8Linear is None: + pytest.skip("torchao MXFP8Linear is unavailable") + from torchtitan.distributed.utils import get_spmd_backend, set_spmd_backend + + previous_backend = get_spmd_backend() + set_spmd_backend("partial_dtensor") + try: + with pytest.raises(ValueError, match="spmd_backend"): + MXFP8Linear.Config(in_features=128, out_features=128).build() + finally: + set_spmd_backend(previous_backend) + + +def test_mxfp8_converter_replaces_a_root_linear_config(monkeypatch): + """A Linear.Config with no parent is returned, not mutated in place. + + ``convert`` writes into ``parent`` for nested configs, so the root case is + the one branch that has to return the replacement. Not covered by the FQN + test below, which passes a FeedForward and so always has a parent. + """ + import torchtitan.components.quantization.mxfp8.converter as converter_mod + + monkeypatch.setattr(converter_mod, "has_cuda_capability", lambda *_: True) + converter = MXFP8LinearConverter( + MXFP8LinearConverter.Config( + model_compile_enabled=True, + ) + ) + + converted = converter.convert( + Linear.Config(in_features=128, out_features=128, bias=False) + ) + + assert isinstance(converted, MXFP8Linear.Config) + assert converted.input_activation_format_for_backward == "bf16" + + +def test_mxfp8_converter_applies_mxfp8_saved_input_fqns(monkeypatch): + import torchtitan.components.quantization.mxfp8.converter as converter_mod + + monkeypatch.setattr(converter_mod, "has_cuda_capability", lambda *_: True) + converter = MXFP8LinearConverter( + MXFP8LinearConverter.Config( + model_compile_enabled=True, + linears_saving_inputs_for_backward_in_mxfp8=["w2"], + ) + ) + converted = converter.convert( + FeedForward.Config( + w1=Linear.Config(in_features=128, out_features=128), + w2=Linear.Config(in_features=128, out_features=128), + w3=Linear.Config(in_features=128, out_features=128), + ) + ) + + assert isinstance(converted.w1, MXFP8Linear.Config) + assert isinstance(converted.w2, MXFP8Linear.Config) + assert isinstance(converted.w3, MXFP8Linear.Config) + assert converted.w1.input_activation_format_for_backward == "bf16" + assert converted.w2.input_activation_format_for_backward == "mxfp8" + assert converted.w3.input_activation_format_for_backward == "bf16" + + +def test_mxfp8_converter_rejects_unmatched_saved_input_fqns(monkeypatch): + import torchtitan.components.quantization.mxfp8.converter as converter_mod + + monkeypatch.setattr(converter_mod, "has_cuda_capability", lambda *_: True) + converter = MXFP8LinearConverter( + MXFP8LinearConverter.Config( + model_compile_enabled=True, + linears_saving_inputs_for_backward_in_mxfp8=["missing"], + ) + ) + model_config = FeedForward.Config( + w1=Linear.Config(in_features=128, out_features=128), + w2=Linear.Config(in_features=128, out_features=128), + w3=Linear.Config(in_features=128, out_features=128), + ) + + with pytest.raises( + ValueError, + match="selectors did not match any converted Linear.Config", + ): + converter.convert(model_config) + + +def test_mxfp8_converter_rejects_empty_saved_input_fqn(): + with pytest.raises(ValueError, match="cannot contain an empty FQN selector"): + MXFP8LinearConverter.Config( + model_compile_enabled=True, + linears_saving_inputs_for_backward_in_mxfp8=[""], + ) + + +@pytest.mark.parametrize( + "config_factory, mxfp8_fqns", + [ + ( + "llama3", + ("attention.qkv_linear.wqkv", "feed_forward.w2"), + ), + ( + "llama3_graph", + ("attention.qkv_linear.wqkv", "feed_forward.w2"), + ), + ( + "deepseek_v3", + ("attention.wkv_b", "feed_forward.w2", "shared_experts.w2"), + ), + ( + "deepseek_v3_graph", + ("attention.wkv_b", "feed_forward.w2", "shared_experts.w2"), + ), + ], +) +def test_builtin_mxfp8_configs_assign_input_activation_format_for_backward( + monkeypatch, config_factory, mxfp8_fqns +): + if MXFP8Linear is None: + pytest.skip("torchao MXFP8Linear is unavailable") + import torchtitan.components.quantization.mxfp8.converter as converter_mod + + monkeypatch.setattr(converter_mod, "has_cuda_capability", lambda *_: True) + if config_factory == "llama3": + from torchtitan.models.llama3.config_registry import ( + llama3_debugmodel_mxfp8 as build_config, + ) + elif config_factory == "llama3_graph": + from torchtitan.experiments.graph_trainer.llama3.config_registry import ( + graph_trainer_llama3_debugmodel_mxfp8 as build_config, + ) + elif config_factory == "deepseek_v3": + from torchtitan.models.deepseek_v3.config_registry import ( + deepseek_v3_debugmodel_mxfp8 as build_config, + ) + else: + from torchtitan.experiments.graph_trainer.deepseek_v3.config_registry import ( + graph_trainer_deepseek_v3_debugmodel_mxfp8 as build_config, + ) + + trainer_config = build_config() + assert trainer_config.model_spec is not None + model_config = trainer_config.model_spec.model + assignments = { + fqn: config.input_activation_format_for_backward + for fqn, config, _parent, _attr in model_config.traverse(MXFP8Linear.Config) + } + + assert assignments + assert "bf16" in assignments.values() + assert "mxfp8" in assignments.values() + for fqn, save_format in assignments.items(): + expected = ( + "mxfp8" if any(selector in fqn for selector in mxfp8_fqns) else "bf16" + ) + assert save_format == expected, f"Unexpected policy for {fqn}" + + +def test_mxfp8_linear_loads_stock_checkpoint(): + pytest.importorskip("torchao") + if MXFP8Linear is None: + pytest.skip("torchao MXFP8Linear is unavailable") + from torchtitan.components.quantization.mxfp8.tensor import ( + _LinearShardedTensorWithMXFP8Compute, + ) + + stock = Linear.Config(in_features=128, out_features=96).build() + mxfp8 = MXFP8Linear.Config(in_features=128, out_features=96).build() + with torch.no_grad(): + stock.weight.normal_() + + mxfp8.load_state_dict(stock.state_dict()) + assert isinstance(mxfp8.weight, _LinearShardedTensorWithMXFP8Compute) + assert torch.equal(mxfp8.weight._tensor, stock.weight) diff --git a/tests/unit_tests/cpu/test_torch_checkpointing.py b/tests/unit_tests/cpu/test_torch_checkpointing.py index e8fa2d85a0..99b788e7b0 100644 --- a/tests/unit_tests/cpu/test_torch_checkpointing.py +++ b/tests/unit_tests/cpu/test_torch_checkpointing.py @@ -6,24 +6,35 @@ import dataclasses import json +import queue import unittest +from concurrent.futures import Future +from contextlib import nullcontext +from pathlib import Path from unittest import mock +import torch import torch.nn as nn + +from torch.distributed.checkpoint.stateful import Stateful from torch_checkpointing.barriers import TCPStoreBarrierConfig from torch_checkpointing.checkpoint_manager import ( CheckpointManager as BackendCheckpointManager, ) -from torch_checkpointing.config import AsyncCheckpointSaverConfig +from torch_checkpointing.config import ( + AsyncCheckpointSaverConfig, + SyncCheckpointSaverConfig, +) from torch_checkpointing.default_resharder import DefaultResharder - from torchtitan.components.checkpointer import ( BaseCheckpointManager, CheckpointManager, + CheckpointStorage, MODEL, OPTIMIZER, ) from torchtitan.components.checkpointer.torch_checkpointing import ( + _async_save_config, _default_backend_config, DEFAULT_TORCH_CHECKPOINTING_BARRIER_TCPSTORE_PORT, TorchCheckpointingManager, @@ -34,14 +45,43 @@ class _BackendManager: def __init__(self) -> None: self.closed = False + self.lock_calls = 0 + self.prewarm_calls = [] + self.save_calls = [] + self.save_result = Future() + + def save(self, checkpoint_id, checkpoint): + self.save_calls.append((checkpoint_id, checkpoint)) + return self.save_result + + def prewarm_staging(self, checkpoint) -> None: + self.prewarm_calls.append(checkpoint) + + def lock(self): + self.lock_calls += 1 + return nullcontext() def close(self) -> None: self.closed = True +class _Stateful(Stateful): + def __init__(self, value) -> None: + self.value = value + + def state_dict(self): + return {"value": self.value} + + def load_state_dict(self, state_dict) -> None: + self.value = state_dict["value"] + + class TorchCheckpointingManagerTest(unittest.TestCase): def _build_manager( - self, config: TorchCheckpointingManager.Config + self, + config: TorchCheckpointingManager.Config, + *, + storage_config=None, ) -> tuple[TorchCheckpointingManager, _BackendManager]: backend_manager = _BackendManager() with mock.patch.object( @@ -52,11 +92,12 @@ def _build_manager( manager = config.build( dataloader=None, model_parts=[nn.Linear(2, 2)], - optimizers=object(), - lr_schedulers=object(), - states={}, + optimizers=_Stateful("optimizer"), + lr_schedulers=_Stateful("scheduler"), + states={"train_state": _Stateful("train")}, sd_adapter=None, base_folder="/tmp", + storage_config=storage_config, ) return manager, backend_manager @@ -105,9 +146,17 @@ def test_disabled_manager_lifecycle_is_noop(self) -> None: self.assertIsNone(manager.maybe_wait_for_staging()) manager.close() + def test_del_ignores_manager_whose_construction_failed(self) -> None: + manager = TorchCheckpointingManager.__new__(TorchCheckpointingManager) + manager.enable = True + manager.save_future = None + manager.purge_thread = None + + manager.__del__() + @mock.patch.dict("os.environ", {"MASTER_ADDR": "checkpoint-host"}) def test_default_backend_configuration_owns_schema_and_barrier(self) -> None: - backend_config = _default_backend_config() + backend_config = _default_backend_config(_async_save_config()) self.assertIsInstance(backend_config.save, AsyncCheckpointSaverConfig) self.assertTrue(backend_config.save.staging_config.use_pinned_memory) @@ -148,8 +197,348 @@ def test_remote_checkpoint_paths_are_rejected_at_construction(self) -> None: with self.assertRaisesRegex(ValueError, rf"{field}.*not yet supported"): self._build_manager(config) + def test_storage_config_is_an_init_parameter_not_a_config_field(self) -> None: + # Backend storage is passed programmatically rather than declared on + # Config, which is Tyro-parsed and not the place for a storage object. + field_names = { + field.name for field in dataclasses.fields(TorchCheckpointingManager.Config) + } + self.assertNotIn("storage_config", field_names) + + storage = mock.Mock() + storage_config = mock.Mock() + storage_config.create_storage.return_value = storage + config = TorchCheckpointingManager.Config( + enable=True, + keep_latest_k=0, + initial_load_model_only=False, + load_only=True, + ) + + manager, _ = self._build_manager(config, storage_config=storage_config) + + # Path probes go through the same storage, wrapped to answer the + # CheckpointStorage protocol, and it is pushed into the backend config + # so saves and loads use it too. + self.assertIsInstance(manager._storage, CheckpointStorage) + manager._storage.isdir("/somewhere") + storage.isdir.assert_called_once_with(Path("/somewhere")) + self.assertIs(storage_config, manager._manager_config.storage_config) + manager.close() + def test_legacy_config_has_no_backend_selector(self) -> None: config = CheckpointManager.Config() self.assertFalse(hasattr(config, "save_backend")) self.assertFalse(hasattr(config, "load_backend")) + + def test_save_obeys_cadence_and_tracks_backend_future(self) -> None: + config = TorchCheckpointingManager.Config( + enable=True, + interval=3, + keep_latest_k=0, + initial_load_model_only=False, + ) + manager, backend_manager = self._build_manager(config) + + self.assertFalse(manager.save(curr_step=1)) + self.assertEqual([], backend_manager.save_calls) + + self.assertTrue(manager.save(curr_step=3)) + self.assertEqual(1, len(backend_manager.save_calls)) + checkpoint_id, checkpoint = backend_manager.save_calls[0] + self.assertEqual("/tmp/checkpoint/step-3", checkpoint_id) + self.assertEqual("train", checkpoint["train_state"]["value"]) + self.assertIs(backend_manager.save_result, manager.save_future) + + backend_manager.save_result.set_result(None) + manager.maybe_wait_for_saving() + self.assertIsNone(manager.save_future) + manager.close() + + def test_prewarm_runs_once_before_first_scheduled_save(self) -> None: + config = TorchCheckpointingManager.Config( + enable=True, + interval=10, + keep_latest_k=0, + initial_load_model_only=False, + ) + manager, backend_manager = self._build_manager(config) + + self.assertFalse(manager.save(curr_step=1)) + self.assertFalse(manager.save(curr_step=2)) + + self.assertEqual(1, len(backend_manager.prewarm_calls)) + self.assertEqual(set(manager.states), set(backend_manager.prewarm_calls[0])) + manager.close() + + def test_load_only_uses_synchronous_backend(self) -> None: + config = TorchCheckpointingManager.Config( + enable=True, + keep_latest_k=0, + initial_load_model_only=False, + load_only=True, + ) + + manager, _ = self._build_manager(config) + + self.assertIsInstance(manager._manager_config.save, SyncCheckpointSaverConfig) + self.assertIsNone(manager._manager_config.save.writer_config.barrier_config) + manager.close() + + def test_load_only_does_not_construct_checkpoint_barrier(self) -> None: + config = TorchCheckpointingManager.Config( + enable=True, + keep_latest_k=0, + initial_load_model_only=False, + load_only=True, + ) + + with mock.patch.object( + TCPStoreBarrierConfig, + "create_barrier", + side_effect=AssertionError("checkpoint barrier constructed"), + ): + manager = config.build( + dataloader=None, + model_parts=[nn.Linear(2, 2)], + optimizers=_Stateful("optimizer"), + lr_schedulers=_Stateful("scheduler"), + states={"train_state": _Stateful("train")}, + sd_adapter=None, + base_folder="/tmp", + ) + + manager.close() + + def test_staging_wait_uses_backend_lock(self) -> None: + config = TorchCheckpointingManager.Config( + enable=True, + keep_latest_k=0, + initial_load_model_only=False, + ) + manager, backend_manager = self._build_manager(config) + + manager.maybe_wait_for_staging() + + self.assertEqual(1, backend_manager.lock_calls) + manager.close() + + def test_save_wait_uses_configured_timeout(self) -> None: + config = TorchCheckpointingManager.Config( + enable=True, + keep_latest_k=0, + initial_load_model_only=False, + ) + manager, _ = self._build_manager(config) + save_future = mock.Mock() + manager.save_future = save_future + + manager.maybe_wait_for_saving() + + save_future.result.assert_called_once_with( + timeout=manager._manager_config.save.wait_timeout_secs + ) + self.assertIsNone(manager.save_future) + manager.close() + + def test_close_releases_resources_when_save_fails(self) -> None: + config = TorchCheckpointingManager.Config( + enable=True, + keep_latest_k=2, + initial_load_model_only=False, + ) + manager, backend_manager = self._build_manager(config) + backend_manager.save_result.set_exception(RuntimeError("save failed")) + manager.save_future = backend_manager.save_result + + with self.assertRaisesRegex(RuntimeError, "save failed"): + manager.close() + + self.assertTrue(backend_manager.closed) + self.assertFalse(manager.purge_thread.is_alive()) + self.assertIsNone(manager.save_future) + + def test_purge_runs_before_this_step_save_is_issued(self) -> None: + # Purging while a save is in flight would let it see, and delete, that + # save's own temporary directory. + config = TorchCheckpointingManager.Config( + enable=True, + interval=1, + keep_latest_k=0, + initial_load_model_only=False, + ) + manager, backend_manager = self._build_manager(config) + calls = [] + backend_manager.save = lambda checkpoint_id, checkpoint: ( + calls.append("save"), + backend_manager.save_result, + )[1] + + with mock.patch.object( + manager, + "_purge_stale_checkpoints", + side_effect=lambda *, saving_step, staging_dir_prefix=None: calls.append( + ("purge", saving_step, staging_dir_prefix) + ), + ): + self.assertTrue(manager.save(curr_step=1)) + + self.assertEqual( + [ + ( + "purge", + 1, + manager._manager_config.save.writer_config.temp_dir_prefix, + ), + "save", + ], + calls, + ) + backend_manager.save_result.set_result(None) + manager.close() + + def _purge_manager(self, keep_latest_k: int, entries: list[str]): + manager = TorchCheckpointingManager.__new__(TorchCheckpointingManager) + manager.keep_latest_k = keep_latest_k + manager.folder = "/checkpoint" + manager.purge_queue = queue.Queue() + manager.purge_thread = object() + manager._storage = mock.Mock(spec=CheckpointStorage) + manager._storage.isdir.return_value = True + manager._storage.listdir.return_value = entries + manager._storage.isfile.return_value = True + return manager + + def _purge(self, manager, *, saving_step: int) -> set[str]: + manager._purge_stale_checkpoints( + saving_step=saving_step, + staging_dir_prefix="tmp_", + ) + purged = set() + while not manager.purge_queue.empty(): + purged.add(manager.purge_queue.get_nowait()) + return purged + + @mock.patch("torch.distributed.get_rank", return_value=0) + def test_purge_reserves_a_slot_for_the_imminent_save(self, _rank) -> None: + # keep_latest_k=3 with a save about to be issued leaves 2 on disk. + manager = self._purge_manager(3, ["step-1", "step-2", "step-3"]) + + self.assertEqual({"/checkpoint/step-1"}, self._purge(manager, saving_step=4)) + + @mock.patch("torch.distributed.get_rank", return_value=0) + def test_purge_treats_metadata_less_directories_as_abandoned(self, _rank) -> None: + # A step-N other than the current save with no metadata is the residue + # of a save that died. It is removed instead of occupying a retained slot. + manager = self._purge_manager(2, ["step-1", "step-2", "step-3"]) + manager._storage.isfile.return_value = False + + self.assertEqual( + { + "/checkpoint/step-1", + "/checkpoint/step-2", + "/checkpoint/step-3", + }, + self._purge(manager, saving_step=4), + ) + manager._storage.remove.assert_not_called() + + @mock.patch("torch.distributed.get_rank", return_value=0) + def test_purge_deletes_abandoned_temporaries_without_spending_a_slot( + self, _rank + ) -> None: + # Two saves failed before their rename. They must not push step-1, the + # only checkpoint this run can resume from, out of the retained set. + manager = self._purge_manager(3, ["step-1", "tmp_step-2", "tmp_step-3"]) + + # step-1 survives, and the purge thread removes the temporaries. + self.assertEqual( + {"/checkpoint/tmp_step-2", "/checkpoint/tmp_step-3"}, + self._purge(manager, saving_step=4), + ) + manager._storage.remove.assert_not_called() + + @mock.patch("torch.distributed.get_rank", return_value=0) + def test_purge_preserves_the_current_staging_directory(self, _rank) -> None: + manager = self._purge_manager(2, ["tmp_step-2", "tmp_step-3"]) + + self.assertEqual( + {"/checkpoint/tmp_step-2"}, self._purge(manager, saving_step=3) + ) + manager._storage.remove.assert_not_called() + + @mock.patch("torch.distributed.get_rank", return_value=0) + def test_purge_ignores_directories_the_backend_did_not_write(self, _rank) -> None: + manager = self._purge_manager( + 2, ["step-1", "step-2", "step-3.partial", "step-4-notes", "notes-step-5"] + ) + + self.assertEqual({"/checkpoint/step-1"}, self._purge(manager, saving_step=6)) + + def test_last_step_export_converts_only_mismatched_float_tensors(self) -> None: + # Regression for the two defects #4166 fixed on the DCP side: gating on + # export_dtype != float32 skipped BF16-to-FP32 exports entirely, and the + # blanket cast turned integer and boolean buffers into floats. + cases = (("float32", torch.float32), ("bfloat16", torch.bfloat16)) + for export_dtype, expected in cases: + with self.subTest(export_dtype=export_dtype): + config = TorchCheckpointingManager.Config( + enable=True, + keep_latest_k=0, + initial_load_model_only=False, + last_save_model_only=True, + export_dtype=export_dtype, + ) + manager, _ = self._build_manager(config) + sync_manager = _BackendManager() + sync_manager.save_result = None + model = _Stateful(None) + model.state_dict = lambda: { + "weight": torch.ones(2, dtype=torch.bfloat16), + "step_count": torch.ones(2, dtype=torch.int64), + "mask": torch.ones(2, dtype=torch.bool), + } + manager.states[MODEL] = model + + with mock.patch.object( + BackendCheckpointManager.Config, + "build", + return_value=sync_manager, + ): + self.assertTrue(manager.save(curr_step=5, last_step=True)) + + saved = sync_manager.save_calls[0][1][MODEL] + self.assertEqual(expected, saved["weight"].dtype) + self.assertEqual(torch.int64, saved["step_count"].dtype) + self.assertEqual(torch.bool, saved["mask"].dtype) + manager.close() + + def test_last_step_uses_synchronous_manager_and_model_only_payload(self) -> None: + config = TorchCheckpointingManager.Config( + enable=True, + keep_latest_k=0, + initial_load_model_only=False, + last_save_model_only=True, + ) + manager, backend_manager = self._build_manager(config) + sync_manager = _BackendManager() + sync_manager.save_result = None + + with mock.patch.object( + BackendCheckpointManager.Config, + "build", + autospec=True, + return_value=sync_manager, + ) as build: + self.assertTrue(manager.save(curr_step=5, last_step=True)) + + self.assertTrue(backend_manager.closed) + self.assertEqual(1, len(sync_manager.save_calls)) + checkpoint_id, checkpoint = sync_manager.save_calls[0] + self.assertEqual("/tmp/checkpoint/step-5", checkpoint_id) + self.assertEqual({MODEL}, set(checkpoint)) + self.assertTrue(sync_manager.closed) + sync_config = build.call_args.args[0] + self.assertIsNone(sync_config.pre_finalize_callback) + manager.close() diff --git a/tests/unit_tests/cpu/test_triton_offset_rmsnorm_override.py b/tests/unit_tests/cpu/test_triton_offset_rmsnorm_override.py new file mode 100644 index 0000000000..b2704ee3a8 --- /dev/null +++ b/tests/unit_tests/cpu/test_triton_offset_rmsnorm_override.py @@ -0,0 +1,108 @@ +# 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 spmd_types as spmd +import torch + +from torchtitan.config import apply_overrides, OverrideConfig +from torchtitan.models.common.decoder_sharding import dense_param_placement +from torchtitan.models.qwen3_5 import model_registry +from torchtitan.models.qwen3_5.model import OffsetRMSNorm +from torchtitan.overrides.offset_rmsnorm import ( + triton_offset_rmsnorm, + TritonOffsetRMSNorm, +) +from torchtitan.protocols.sharding import ShardingConfig + + +class TestTritonOffsetRMSNormOverride(unittest.TestCase): + def test_override_replaces_all_qwen35_offset_norms(self): + config = model_registry("debugmodel", attn_backend="flex").model + num_offset_norms = len(list(config.traverse(OffsetRMSNorm.Config))) + + replacements = apply_overrides( + OverrideConfig( + imports=["torchtitan.overrides.offset_rmsnorm." "triton_offset_rmsnorm"] + ), + config, + ) + + self.assertGreater(num_offset_norms, 0) + self.assertEqual(len(replacements), num_offset_norms) + self.assertEqual( + len(list(config.traverse(TritonOffsetRMSNorm.Config))), + num_offset_norms, + ) + + def test_config_is_replaced_without_changing_state_dict(self): + stock_config = OffsetRMSNorm.Config( + dim=32, + eps=1e-5, + param_init={"weight": torch.nn.init.zeros_}, + ) + + replacement = triton_offset_rmsnorm(stock_config) + + self.assertIsInstance(replacement, TritonOffsetRMSNorm.Config) + self.assertEqual(replacement.dim, stock_config.dim) + self.assertEqual(replacement.eps, stock_config.eps) + self.assertIs(replacement.param_init, stock_config.param_init) + self.assertEqual( + list(replacement.build().state_dict()), + list(stock_config.build().state_dict()), + ) + + def test_override_adds_local_compute_region_for_sharded_norm(self): + activation = spmd.SpmdType( + {"dp": spmd.V, "tp": spmd.I}, + partition_spec=spmd.PartitionSpec("dp", None), + ) + weight = dense_param_placement(tp=spmd.R) + sharding = ShardingConfig( + state_shardings={"weight": weight}, + in_src_shardings={"input": activation}, + out_src_shardings=activation, + ) + stock_config = OffsetRMSNorm.Config( + dim=32, + sharding_config=sharding, + ) + + replacement = triton_offset_rmsnorm(stock_config) + + self.assertIsNotNone(replacement.sharding_config) + assert replacement.sharding_config is not None + self.assertIsNotNone(replacement.sharding_config.local_map) + assert replacement.sharding_config.local_map is not None + self.assertEqual( + replacement.sharding_config.local_map.in_grad_placements, + (activation,), + ) + self.assertIsNotNone(replacement.weight_grad_sharding) + assert replacement.weight_grad_sharding is not None + self.assertEqual( + replacement.weight_grad_sharding.local_type["tp"], + spmd.P, + ) + + def test_cpu_fallback_matches_stock_module(self): + config = OffsetRMSNorm.Config(dim=32, eps=1e-6) + stock = config.build() + fused = triton_offset_rmsnorm(config).build() + with torch.no_grad(): + weight = torch.randn(32) + stock.weight.copy_(weight) + fused.weight.copy_(weight) + + input = torch.randn(4, 32) + + torch.testing.assert_close(fused(input), stock(input)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit_tests/gpu/test_mxfp8_fsdp.py b/tests/unit_tests/gpu/test_mxfp8_fsdp.py new file mode 100644 index 0000000000..aa328a0632 --- /dev/null +++ b/tests/unit_tests/gpu/test_mxfp8_fsdp.py @@ -0,0 +1,544 @@ +# 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 os + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.elastic.utils.distributed import get_free_port +from torch.distributed.fsdp import fully_shard, MixedPrecisionPolicy +from torch.distributed.tensor import DTensor + + +pytest.importorskip("torchao") +pytest.importorskip("torchao.prototype.moe_training.kernels.mxfp8") + +import torchtitan.components.quantization.mxfp8.tensor as mxfp8_tensor # noqa: E402 +from torchtitan.components.quantization._fsdp_tensor import ( # noqa: E402 + _UnshardedFSDPTensor, +) +from torchtitan.components.quantization.mxfp8.linear import MXFP8Linear # noqa: E402 +from torchtitan.components.quantization.mxfp8.tensor import ( # noqa: E402 + _LinearShardedTensorWithMXFP8Compute, +) +from torchtitan.distributed.cudagraph import ( # noqa: E402 + cudagraph_teardown, + CUDAGraphWrapper, +) +from torchtitan.experiments.graph_trainer.simple_fsdp import ( # noqa: E402 + data_parallel, + disable_active_parametrization, + MixedPrecisionPolicy as SimpleFSDPMixedPrecisionPolicy, +) + + +# Every test here spawns a two-rank process group, so the whole module belongs +# to the multi_gpu lane. Without the marker the tests land in the single-GPU +# lane instead, where the device-count guard skips all of them. +pytestmark = [ + pytest.mark.multi_gpu, + pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires two GPUs"), + pytest.mark.skipif( + torch.cuda.is_available() and torch.cuda.get_device_capability() < (10, 0), + reason="MXFP8 requires SM100 or later", + ), +] + + +def _get_weight_param(linear): + state = fully_shard.state(linear) + param_group = state._fsdp_param_group + assert param_group is not None + return next( + param + for param in param_group.fsdp_params + if param._module_info.param_name == "weight" + ) + + +def _run_reshard_after_forward( + rank: int, + world_size: int, + port: int, +) -> None: + """Test RAF=true release and refill of FSDP-managed MXFP8 operands. + + Forward and backward use separate unshards. Each reshard must release both + the temporary BF16 all-gather output and the MXFP8 operand storage, while a + later unshard must refill the same stable inner tensor objects. + """ + os.environ["MASTER_ADDR"] = "localhost" + os.environ["MASTER_PORT"] = str(port) + torch.cuda.set_device(rank) + dist.init_process_group("nccl", rank=rank, world_size=world_size) + try: + mesh = init_device_mesh("cuda", (world_size,), mesh_dim_names=("dp_shard",)) + linear = ( + MXFP8Linear.Config( + in_features=128, + out_features=128, + bias=False, + ) + .build() + .cuda() + .bfloat16() + ) + linear.compile() + fully_shard( + linear, + mesh=mesh, + mp_policy=MixedPrecisionPolicy( + param_dtype=torch.bfloat16, + reduce_dtype=torch.bfloat16, + ), + reshard_after_forward=True, + ) + assert isinstance( + linear.weight.to_local(), _LinearShardedTensorWithMXFP8Compute + ) + + input_MK = torch.randn( + 64, + 128, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + output_MN = linear(input_MK) + weight_param = _get_weight_param(linear) + inner_tensor_ids = tuple(map(id, weight_param._unsharded_inner_tensors)) + assert isinstance( + linear.weight.to_local(), _LinearShardedTensorWithMXFP8Compute + ) + assert all( + tensor.untyped_storage().size() == 0 + for tensor in weight_param.all_gather_outputs + ) + assert all( + tensor.untyped_storage().size() == 0 + for tensor in weight_param._unsharded_inner_tensors + ) + + output_MN.sum().backward() + assert isinstance( + linear.weight.to_local(), _LinearShardedTensorWithMXFP8Compute + ) + assert tuple(map(id, weight_param._unsharded_inner_tensors)) == inner_tensor_ids + assert all( + tensor.untyped_storage().size() == 0 + for tensor in weight_param.all_gather_outputs + ) + assert all( + tensor.untyped_storage().size() == 0 + for tensor in weight_param._unsharded_inner_tensors + ) + finally: + dist.destroy_process_group() + + +def _run_pp_cache_lifecycle( + rank: int, + world_size: int, + port: int, +) -> None: + """Test RAF=false cache reuse across pipeline-parallel microbatches. + + The first unshard quantizes the weight once. Multiple forwards and + backwards reuse those MXFP8 operands until the last backward requests + a reshard. The next generation must quantize again while reusing the same + unsharded inner tensor objects. + """ + os.environ["MASTER_ADDR"] = "localhost" + os.environ["MASTER_PORT"] = str(port) + torch.cuda.set_device(rank) + dist.init_process_group("nccl", rank=rank, world_size=world_size) + original_quantize_weight = mxfp8_tensor._quantize_mxfp8_weight + num_quantize_calls = 0 + + def counted_quantize_weight(weight_NK: torch.Tensor): + nonlocal num_quantize_calls + num_quantize_calls += 1 + return original_quantize_weight(weight_NK) + + mxfp8_tensor._quantize_mxfp8_weight = counted_quantize_weight + try: + mesh = init_device_mesh("cuda", (world_size,), mesh_dim_names=("dp_shard",)) + linear = ( + MXFP8Linear.Config( + in_features=128, + out_features=128, + bias=False, + ) + .build() + .cuda() + .bfloat16() + ) + fully_shard( + linear, + mesh=mesh, + mp_policy=MixedPrecisionPolicy( + param_dtype=torch.bfloat16, + reduce_dtype=torch.bfloat16, + ), + reshard_after_forward=False, + ) + linear.set_is_last_backward(False) + linear.set_reshard_after_backward(False) + linear.set_requires_gradient_sync(False) + + inputs = [ + torch.randn( + 64, + 128, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + for _ in range(2) + ] + outputs = [linear(input_MK) for input_MK in inputs] + assert num_quantize_calls == 1 + weight_param = _get_weight_param(linear) + assert isinstance(linear.weight, _UnshardedFSDPTensor) + assert linear.weight.operands is not None + assert len(weight_param._unsharded_inner_tensors) == 3 + operands = linear.weight.operands + assert operands is not None + assert ( + operands.weight_qdata_dgrad_NK.data_ptr() + == operands.weight_qdata_fprop_KN.data_ptr() + ) + inner_tensor_ids = tuple(map(id, weight_param._unsharded_inner_tensors)) + assert all( + tensor.untyped_storage().size() == 0 + for tensor in weight_param.all_gather_outputs + ) + assert all( + tensor.untyped_storage().size() > 0 + for tensor in weight_param._unsharded_inner_tensors + ) + + outputs[0].sum().backward() + assert num_quantize_calls == 1 + assert isinstance(linear.weight, _UnshardedFSDPTensor) + assert linear.weight.operands is not None + assert all( + tensor.untyped_storage().size() == 0 + for tensor in weight_param.all_gather_outputs + ) + assert all( + tensor.untyped_storage().size() > 0 + for tensor in weight_param._unsharded_inner_tensors + ) + + linear.set_is_last_backward(True) + linear.set_reshard_after_backward(True) + linear.set_requires_gradient_sync(True) + outputs[1].sum().backward() + assert num_quantize_calls == 1 + assert isinstance( + linear.weight.to_local(), _LinearShardedTensorWithMXFP8Compute + ) + assert all( + tensor.untyped_storage().size() == 0 + for tensor in weight_param.all_gather_outputs + ) + assert all( + tensor.untyped_storage().size() == 0 + for tensor in weight_param._unsharded_inner_tensors + ) + + output_MN = linear(inputs[0].detach()) + assert num_quantize_calls == 2 + assert isinstance(linear.weight, _UnshardedFSDPTensor) + assert linear.weight.operands is not None + assert tuple(map(id, weight_param._unsharded_inner_tensors)) == inner_tensor_ids + assert all( + tensor.untyped_storage().size() == 0 + for tensor in weight_param.all_gather_outputs + ) + assert all( + tensor.untyped_storage().size() > 0 + for tensor in weight_param._unsharded_inner_tensors + ) + output_MN.sum().backward() + finally: + mxfp8_tensor._quantize_mxfp8_weight = original_quantize_weight + dist.destroy_process_group() + + +def _run_cuda_graph_cache_lifecycle( + rank: int, + world_size: int, + port: int, +) -> None: + """Test that the RAF=false MXFP8 cache is safe for CUDA graph replay. + + Warmup, capture, and replay must not re-quantize the weight or change the + cached operand addresses. An explicit reshard after graph teardown must + release the FSDP-managed operand storage. + """ + os.environ["MASTER_ADDR"] = "localhost" + os.environ["MASTER_PORT"] = str(port) + torch.cuda.set_device(rank) + dist.init_process_group("nccl", rank=rank, world_size=world_size) + original_quantize_weight = mxfp8_tensor._quantize_mxfp8_weight + num_quantize_calls = 0 + + def counted_quantize_weight(weight_NK: torch.Tensor): + nonlocal num_quantize_calls + num_quantize_calls += 1 + return original_quantize_weight(weight_NK) + + mxfp8_tensor._quantize_mxfp8_weight = counted_quantize_weight + try: + mesh = init_device_mesh("cuda", (world_size,), mesh_dim_names=("dp_shard",)) + linear = ( + MXFP8Linear.Config( + in_features=128, + out_features=128, + bias=False, + ) + .build() + .cuda() + .bfloat16() + ) + fully_shard( + linear, + mesh=mesh, + mp_policy=MixedPrecisionPolicy( + param_dtype=torch.bfloat16, + reduce_dtype=torch.bfloat16, + ), + reshard_after_forward=False, + ) + linear.set_is_last_backward(False) + linear.set_reshard_after_backward(False) + linear.set_requires_gradient_sync(False) + + def forward_backward( + input_MK: torch.Tensor, + ) -> torch.Tensor: + output_MN = linear(input_MK) + output_MN.sum().backward() + return output_MN + + input_MK = torch.randn( + 64, + 128, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + + # Establish the FSDP unsharded generation and its prepared weights on + # the current stream before CUDA-graph warmup moves to its side stream. + forward_backward(input_MK) + torch.cuda.synchronize() + assert num_quantize_calls == 1 + weight_param = _get_weight_param(linear) + cache_addresses = tuple( + tensor.data_ptr() for tensor in weight_param._unsharded_inner_tensors + ) + assert all( + tensor.untyped_storage().size() == 0 + for tensor in weight_param.all_gather_outputs + ) + assert all( + tensor.untyped_storage().size() > 0 + for tensor in weight_param._unsharded_inner_tensors + ) + + graphed_step = CUDAGraphWrapper( + forward_backward, + (input_MK,), + static_input_indices=(0,), + should_check_address=True, + ) + + # RAF=false keeps the prepared weights alive, so CUDA-graph warmup, + # capture, and replay reuse the same tensor objects and addresses. + graphed_step(input_MK) + assert num_quantize_calls == 1 + + captured_output_MN = graphed_step(input_MK).clone() + with torch.no_grad(): + input_MK.copy_(torch.randn_like(input_MK)) + replay_output_MN = graphed_step(input_MK).clone() + torch.cuda.synchronize() + + assert graphed_step._graph is not None + assert num_quantize_calls == 1 + assert ( + tuple(tensor.data_ptr() for tensor in weight_param._unsharded_inner_tensors) + == cache_addresses + ) + assert all( + tensor.untyped_storage().size() == 0 + for tensor in weight_param.all_gather_outputs + ) + assert not torch.equal(captured_output_MN, replay_output_MN) + + graphed_step.teardown() + linear.reshard() + assert all( + tensor.untyped_storage().size() == 0 + for tensor in weight_param._unsharded_inner_tensors + ) + finally: + mxfp8_tensor._quantize_mxfp8_weight = original_quantize_weight + cudagraph_teardown() + dist.destroy_process_group() + + +def _run_simple_fsdp( + rank: int, + world_size: int, + port: int, +) -> None: + """Test GraphTrainer SimpleFSDP unsharded tensors and gradient propagation.""" + os.environ["MASTER_ADDR"] = "localhost" + os.environ["MASTER_PORT"] = str(port) + torch.cuda.set_device(rank) + dist.init_process_group("nccl", rank=rank, world_size=world_size) + original_quantize_weight = mxfp8_tensor._quantize_mxfp8_weight + num_quantize_calls = 0 + + def counted_quantize_weight(weight_NK: torch.Tensor): + nonlocal num_quantize_calls + num_quantize_calls += 1 + return original_quantize_weight(weight_NK) + + mxfp8_tensor._quantize_mxfp8_weight = counted_quantize_weight + try: + mesh = init_device_mesh("cuda", (world_size,), mesh_dim_names=("fsdp",)) + linear = ( + MXFP8Linear.Config( + in_features=128, + out_features=128, + bias=False, + ) + .build() + .cuda() + .bfloat16() + ) + linear = data_parallel( + linear, + mesh, + mode="fully_shard", + mp_policy=SimpleFSDPMixedPrecisionPolicy( + param_dtype=torch.bfloat16, + reduce_dtype=torch.bfloat16, + ), + # apply_simple_fsdp() composes this for real GraphTrainer runs. + ) + sharded_weight = linear._parameters["weight"] + assert isinstance( + sharded_weight._local_tensor, _LinearShardedTensorWithMXFP8Compute + ) + + input_MK = torch.randn( + 64, + 128, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + output_MN = linear(input_MK) + output_MN.sum().backward() + + assert output_MN.shape == (64, 128) + assert num_quantize_calls == 1 + assert input_MK.grad is not None + assert sharded_weight.grad is not None + finally: + mxfp8_tensor._quantize_mxfp8_weight = original_quantize_weight + dist.destroy_process_group() + + +@pytest.mark.parametrize( + "target", + [ + _run_reshard_after_forward, + _run_pp_cache_lifecycle, + _run_cuda_graph_cache_lifecycle, + _run_simple_fsdp, + ], + ids=[ + "reshard-after-forward", + "pp-cache", + "cuda-graph-cache", + "simple-fsdp", + ], +) +def test_mxfp8_fsdp_tensor_lifecycle(target): + mp.spawn( + target, + args=(2, get_free_port()), + nprocs=2, + join=True, + ) + + +def _run_simple_fsdp_disabled_parametrization( + rank: int, + world_size: int, + port: int, +) -> None: + """Test that disable_active_parametrization() yields the raw parameter. + + Models call it around ``init_states()`` to inspect and initialize weights. + Building an unsharded tensor there would quantize the still-sharded shard as + if it were the logical tensor, so the disable has to cover that step too. + """ + os.environ["MASTER_ADDR"] = "localhost" + os.environ["MASTER_PORT"] = str(port) + torch.cuda.set_device(rank) + dist.init_process_group("nccl", rank=rank, world_size=world_size) + try: + mesh = init_device_mesh("cuda", (world_size,), mesh_dim_names=("fsdp",)) + linear = ( + MXFP8Linear.Config(in_features=128, out_features=128, bias=False) + .build() + .cuda() + .bfloat16() + ) + linear = data_parallel( + linear, + mesh, + mode="fully_shard", + mp_policy=SimpleFSDPMixedPrecisionPolicy( + param_dtype=torch.bfloat16, + reduce_dtype=torch.bfloat16, + ), + ) + + # Reading the parametrized weight all-gathers, so every rank has to + # reach both of these. + active_weight = linear.weight + with disable_active_parametrization(): + disabled_weight = linear.weight + + assert isinstance(active_weight, _UnshardedFSDPTensor) + assert isinstance(disabled_weight, DTensor) + assert isinstance( + disabled_weight._local_tensor, _LinearShardedTensorWithMXFP8Compute + ) + assert not isinstance(disabled_weight._local_tensor, _UnshardedFSDPTensor) + finally: + dist.destroy_process_group() + + +def test_simple_fsdp_disable_active_parametrization(): + mp.spawn( + _run_simple_fsdp_disabled_parametrization, + args=(2, get_free_port()), + nprocs=2, + join=True, + ) diff --git a/tests/unit_tests/gpu/test_mxfp8_linear.py b/tests/unit_tests/gpu/test_mxfp8_linear.py new file mode 100644 index 0000000000..9066082d6b --- /dev/null +++ b/tests/unit_tests/gpu/test_mxfp8_linear.py @@ -0,0 +1,508 @@ +# 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 pytest +import torch +import torch.nn as nn +from torch.utils.checkpoint import checkpoint + + +pytest.importorskip("torchao") +pytest.importorskip("torchao.prototype.moe_training.kernels.mxfp8") + +import torchtitan.components.quantization.mxfp8.linear as mxfp8_linear # noqa: E402 +from torchtitan.components.quantization._fsdp_tensor import ( # noqa: E402 + _UnshardedFSDPTensor, +) +from torchtitan.components.quantization.mxfp8.linear import MXFP8Linear # noqa: E402 +from torchtitan.components.quantization.mxfp8.tensor import ( # noqa: E402 + _LinearShardedTensorWithMXFP8Compute, +) + + +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required"), + pytest.mark.skipif( + torch.cuda.is_available() and torch.cuda.get_device_capability() < (10, 0), + reason="MXFP8 requires SM100 or later", + ), +] + + +def _make_sharded_mxfp8_linear( + in_features: int = 128, + out_features: int = 96, + *, + bias: bool = True, + input_activation_format_for_backward: str = "bf16", +) -> MXFP8Linear: + """Build a layer in its as-constructed state, before any unshard.""" + return ( + MXFP8Linear.Config( + in_features=in_features, + out_features=out_features, + bias=bias, + input_activation_format_for_backward=input_activation_format_for_backward, + ) + .build() + .cuda() + .bfloat16() + ) + + +def _make_mxfp8_linear( + in_features: int = 128, + out_features: int = 96, + *, + bias: bool = True, + input_activation_format_for_backward: str = "bf16", +) -> MXFP8Linear: + """Build a layer that is ready to run forward. + + ``forward`` requires a data parallel implementation to have built the + unsharded tensor for the current unshard lifetime. These tests are single + process and have none, so stand in for FSDP's post-all-gather hook and + install the unsharded tensor directly. + """ + return _install_unsharded_weight( + _make_sharded_mxfp8_linear( + in_features, + out_features, + bias=bias, + input_activation_format_for_backward=input_activation_format_for_backward, + ) + ) + + +def _build_unsharded_tensor( + sharded_weight: _LinearShardedTensorWithMXFP8Compute, + weight_NK: torch.Tensor, +) -> _UnshardedFSDPTensor: + """Quantize and wrap, as fsdp_post_all_gather does on the first unshard. + + These tests are single process, so the weight is never a DTensor and the + DTensor half of _BuildUnshardedTensorFunction does not apply. + """ + with torch.no_grad(): + return _UnshardedFSDPTensor( + weight_NK, sharded_weight._build_operands(weight_NK) + ) + + +def _install_unsharded_weight(linear: MXFP8Linear) -> MXFP8Linear: + """Stand in for FSDP's post-all-gather hook on a single-process layer. + + The unsharded tensor is storage-free, so anything that reads the weight's + storage -- ``state_dict``, ``load_state_dict`` -- has to run before this. + """ + sharded_weight = linear.weight + linear.weight = nn.Parameter( + _build_unsharded_tensor(sharded_weight, sharded_weight._tensor), + requires_grad=sharded_weight.requires_grad, + ) + return linear + + +@pytest.mark.parametrize("input_activation_format_for_backward", ["bf16", "mxfp8"]) +def test_mxfp8_linear_saves_selected_input_activation( + input_activation_format_for_backward, +): + linear = _make_mxfp8_linear( + input_activation_format_for_backward=input_activation_format_for_backward, + ) + x = torch.randn( + 37, + linear.in_features, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + + saved_tensors = [] + + def pack_hook(tensor): + saved_tensors.append(tensor) + return tensor + + with torch.autograd.graph.saved_tensors_hooks(pack_hook, lambda tensor: tensor): + output = linear(x) + output.backward(torch.randn_like(output)) + + assert output.shape == (37, linear.out_features) + # The weight is saved as the single unsharded-tensor wrapper, not as its + # individual DGRAD operands, so FSDP can free and refill that storage + # around the reshard. Only the activation operands are saved as tensors. + weight_saves = [ + tensor for tensor in saved_tensors if isinstance(tensor, _UnshardedFSDPTensor) + ] + activation_saves = [ + tensor + for tensor in saved_tensors + if not isinstance(tensor, _UnshardedFSDPTensor) + ] + assert len(weight_saves) == 1 + if input_activation_format_for_backward == "bf16": + assert len(activation_saves) == 1 + assert activation_saves[0].dtype == torch.bfloat16 + assert ( + activation_saves[0].untyped_storage()._cdata == x.untyped_storage()._cdata + ) + else: + assert len(activation_saves) == 2 + assert all(tensor.dtype != torch.bfloat16 for tensor in activation_saves) + assert ( + sum(tensor.dtype == torch.float8_e4m3fn for tensor in activation_saves) == 1 + ) + assert ( + sum(tensor.dtype == torch.float8_e8m0fnu for tensor in activation_saves) + == 1 + ) + assert all(type(tensor) is torch.Tensor for tensor in activation_saves) + + +@pytest.mark.parametrize( + ("input_activation_format_for_backward", "expected_quantize_calls"), + [ + ("bf16", [(True, False), (True, True), (False, True)]), + ("mxfp8", [(True, True), (True, True)]), + ], +) +def test_mxfp8_input_activation_format_for_backward_controls_quantization_work( + monkeypatch, + input_activation_format_for_backward, + expected_quantize_calls, +): + original_quantize = mxfp8_linear.mxfp8_quantize_cuda + quantize_calls = [] + + def record_quantize(*args, **kwargs): + quantize_calls.append((kwargs["rowwise"], kwargs["colwise"])) + return original_quantize(*args, **kwargs) + + monkeypatch.setattr(mxfp8_linear, "mxfp8_quantize_cuda", record_quantize) + linear = _make_mxfp8_linear( + bias=False, + input_activation_format_for_backward=input_activation_format_for_backward, + ) + x = torch.randn( + 64, + linear.in_features, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + + linear(x).sum().backward() + + assert quantize_calls == expected_quantize_calls + + +def test_mxfp8_square_weight_dgrad_qdata_is_transpose_view(): + weight_NK = torch.randn( + 96, + 128, + device="cuda", + dtype=torch.bfloat16, + ) + unsharded_tensor = _build_unsharded_tensor( + _LinearShardedTensorWithMXFP8Compute(weight_NK), weight_NK + ) + operands = unsharded_tensor.operands + assert operands is not None + inner_tensor_names, metadata = unsharded_tensor.__tensor_flatten__() + rebuilt_unsharded_tensor = type(unsharded_tensor).__tensor_unflatten__( + {name: getattr(unsharded_tensor, name) for name in inner_tensor_names}, + metadata, + unsharded_tensor.shape, + unsharded_tensor.stride(), + ) + rebuilt_operands = rebuilt_unsharded_tensor.operands + assert rebuilt_operands is not None + + # Inner tensors are named after the operands dataclass fields. + # The FPROP qdata is a property, not a field, so FSDP does not manage it. + assert inner_tensor_names == [ + "_weight_qdata_dgrad_NK", + "_weight_scale_fprop_swizzled", + "_weight_scale_dgrad_swizzled", + ] + assert ( + operands.weight_qdata_dgrad_NK.data_ptr() + == operands.weight_qdata_fprop_KN.data_ptr() + ) + assert torch.equal( + operands.weight_qdata_dgrad_NK, + operands.weight_qdata_fprop_KN.t(), + ) + assert ( + rebuilt_operands.weight_qdata_dgrad_NK.data_ptr() + == rebuilt_operands.weight_qdata_fprop_KN.data_ptr() + ) + + +def test_operands_fields_must_be_distinct_allocations(): + """FSDP owns each field's storage, so a field may not alias another. + + Derived views belong in properties, as ``_MXFP8LinearOperands`` does for + its FPROP qdata. A format that made one a field instead would have FSDP + free the same storage twice. + """ + from dataclasses import dataclass + + from torchtitan.components.quantization._fsdp_tensor import _unsharded_inner_tensors + + qdata = torch.empty(64, 64, device="cuda", dtype=torch.float8_e4m3fn) + + @dataclass(frozen=True) + class AliasingOperands: + qdata_dgrad: torch.Tensor + qdata_fprop: torch.Tensor + + with pytest.raises(ValueError, match="distinct allocations"): + _unsharded_inner_tensors(AliasingOperands(qdata, qdata.t())) + + @dataclass(frozen=True) + class DistinctOperands: + qdata_dgrad: torch.Tensor + scale: torch.Tensor + + scale = torch.empty(64, 2, device="cuda", dtype=torch.float8_e8m0fnu) + assert len(_unsharded_inner_tensors(DistinctOperands(qdata, scale))) == 2 + + +def test_mxfp8_linear_quantizes_per_call_without_an_unsharded_tensor(): + """A layer nobody unsharded builds its operands on every call. + + GraphTrainer under the spmd_types backend reaches forward this way: its + runtime hands over a plain annotated local tensor, so the wrapper + SimpleFSDP built never arrives. Eager FSDP2 always installs one. + """ + linear = _make_sharded_mxfp8_linear() + assert isinstance(linear.weight, _LinearShardedTensorWithMXFP8Compute) + assert not isinstance(linear.weight, _UnshardedFSDPTensor) + x = torch.randn( + 32, + linear.in_features, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + + output = linear(x) + output.sum().backward() + + assert output.shape == (32, linear.out_features) + assert x.grad is not None + # The gradient reaches the wrapped parameter in high precision. + assert linear.weight.grad is not None + assert linear.weight.grad.dtype == torch.bfloat16 + + +def test_mxfp8_linear_gradient_reaches_the_unsharded_tensor(): + linear = _make_mxfp8_linear() + assert isinstance(linear.weight, _UnshardedFSDPTensor) + x = torch.randn( + 32, + linear.in_features, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + + output = linear(x) + output.sum().backward() + + assert output.shape == (32, linear.out_features) + assert x.grad is not None + # The gradient reaches the wrapped parameter in high precision. + assert linear.weight.grad is not None + assert linear.weight.grad.dtype == torch.bfloat16 + + +@pytest.mark.parametrize("input_activation_format_for_backward", ["bf16", "mxfp8"]) +@pytest.mark.parametrize("execution_mode", ["compile", "activation_checkpoint"]) +def test_mxfp8_linear_runs_outside_plain_eager( + execution_mode, + input_activation_format_for_backward, +): + """Both non-eager entry points must reach the weight gradient. + + Each one re-enters forward in a way that can lose the saved activation + state: compile traces it, and non-reentrant checkpointing discards it and + recreates it during recompute. The saved state differs per format, so both + formats are exercised under both. + """ + linear = _make_mxfp8_linear( + bias=False, + input_activation_format_for_backward=input_activation_format_for_backward, + ) + x = torch.randn( + 64, + linear.in_features, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + + if execution_mode == "compile": + output = torch.compile(linear, fullgraph=True)(x) + else: + output = checkpoint(linear, x, use_reentrant=False) + output.backward(torch.randn_like(output)) + + assert output.shape == (64, linear.out_features) + assert x.grad is not None + assert linear.weight.grad is not None + # The gradient is a plain tensor, not the unsharded wrapper. + assert type(linear.weight.grad) is torch.Tensor + + +def test_mxfp8_input_activation_formats_for_backward_match(): + bf16 = _make_sharded_mxfp8_linear( + bias=False, + input_activation_format_for_backward="bf16", + ) + mxfp8 = _make_sharded_mxfp8_linear( + bias=False, + input_activation_format_for_backward="mxfp8", + ) + # Copy the weight while it still has storage, then unshard both layers. + mxfp8.load_state_dict(bf16.state_dict()) + _install_unsharded_weight(bf16) + _install_unsharded_weight(mxfp8) + + x_hp = torch.randn( + 64, + bf16.in_features, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + x_mxfp8 = x_hp.detach().clone().requires_grad_() + grad_output = torch.randn( + 64, + bf16.out_features, + device="cuda", + dtype=torch.bfloat16, + ) + + output_hp = bf16(x_hp) + output_mxfp8 = mxfp8(x_mxfp8) + output_hp.backward(grad_output) + output_mxfp8.backward(grad_output) + + torch.testing.assert_close(output_hp, output_mxfp8, rtol=0, atol=0) + torch.testing.assert_close(x_hp.grad, x_mxfp8.grad, rtol=0, atol=0) + torch.testing.assert_close( + bf16.weight.grad, + mxfp8.weight.grad, + rtol=0, + atol=0, + ) + + +class _StubMesh: + """Stands in for a DeviceMesh: fsdp_pre_all_gather only reads ``size()``.""" + + def __init__(self, size: int) -> None: + self._size = size + + def size(self) -> int: + return self._size + + +class _StubMixedPrecisionPolicy: + param_dtype = torch.bfloat16 + + +def test_fsdp_pre_all_gather_pads_an_uneven_shard(): + """A dim-0 size that does not divide the mesh leaves the last rank short. + + All-gather needs every rank to contribute the same number of elements, so + FSDP's contract is that the hook returns the *padded* shard and passes the + logical size through as metadata. Driven directly rather than through FSDP + because a dense MXFP8 weight cannot be unevenly sharded on two ranks: the + kernels require out_features divisible by 32, which always divides 2. + """ + # Logical (96, 128) over five ranks: ceil(96 / 5) == 20 rows each, so the + # last rank holds only 16 and pads up to 20. + shard_NK = torch.randn(16, 128, device="cuda", dtype=torch.bfloat16) + sharded_weight = _LinearShardedTensorWithMXFP8Compute(shard_NK) + + (comm_NK,), metadata = sharded_weight.fsdp_pre_all_gather( + _StubMesh(5), + torch.Size([96, 128]), + None, + None, + _StubMixedPrecisionPolicy(), + ) + + assert comm_NK.shape == (20, 128) + assert torch.equal(comm_NK[:16], shard_NK) + assert torch.count_nonzero(comm_NK[16:]) == 0 + # The logical size rides along so post-all-gather can drop the padding. + assert tuple(metadata) == (96, 128) + + +def test_fsdp_pre_all_gather_casts_while_padding(): + """A shard needing both a cast and padding must get one buffer, not two. + + The pad path allocates directly in the comm dtype and lets the copy cast, + so this pins that the cast still happens and the real rows survive it. + """ + + class _Fp32Policy: + param_dtype = torch.bfloat16 + reduce_dtype = torch.bfloat16 + + shard_NK = torch.randn(16, 128, device="cuda", dtype=torch.float32) + sharded_weight = _LinearShardedTensorWithMXFP8Compute(shard_NK) + + (comm_NK,), metadata = sharded_weight.fsdp_pre_all_gather( + _StubMesh(5), + torch.Size([96, 128]), + None, + None, + _Fp32Policy(), + ) + + assert comm_NK.dtype == torch.bfloat16 + assert comm_NK.shape == (20, 128) + torch.testing.assert_close(comm_NK[:16], shard_NK.bfloat16()) + assert torch.count_nonzero(comm_NK[16:]) == 0 + assert tuple(metadata) == (96, 128) + + +def test_fsdp_post_all_gather_drops_the_padding(): + """The gathered buffer includes padding; quantization must not see it.""" + sharded_weight = _LinearShardedTensorWithMXFP8Compute( + torch.randn(16, 128, device="cuda", dtype=torch.bfloat16) + ) + # Five ranks contributing 20 padded rows each. + gathered_NK = torch.randn(100, 128, device="cuda", dtype=torch.bfloat16) + + unsharded_tensor, unsharded_inner_tensors = sharded_weight.fsdp_post_all_gather( + (gathered_NK,), torch.Size([96, 128]), torch.bfloat16 + ) + + assert isinstance(unsharded_tensor, _UnshardedFSDPTensor) + assert unsharded_tensor.shape == (96, 128) + assert len(unsharded_inner_tensors) == 3 + + +def test_fsdp_pre_all_gather_rejects_a_non_zero_shard_dim(): + """Only dim 0 is supported; the all-gather concatenates along it.""" + sharded_weight = _LinearShardedTensorWithMXFP8Compute( + torch.randn(96, 64, device="cuda", dtype=torch.bfloat16) + ) + with pytest.raises(NotImplementedError, match="sharding dimension 0 only"): + sharded_weight.fsdp_pre_all_gather( + _StubMesh(2), + torch.Size([96, 128]), + None, + None, + _StubMixedPrecisionPolicy(), + ) diff --git a/tests/unit_tests/gpu/test_triton_offset_rmsnorm_override.py b/tests/unit_tests/gpu/test_triton_offset_rmsnorm_override.py new file mode 100644 index 0000000000..12de488a81 --- /dev/null +++ b/tests/unit_tests/gpu/test_triton_offset_rmsnorm_override.py @@ -0,0 +1,268 @@ +# 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.overrides.offset_rmsnorm import ( + _triton_offset_rms_norm_backward_op, + _triton_offset_rms_norm_op, + triton_offset_rms_norm, +) + + +_EPS = 1e-6 +_FUDGE_FACTOR = 2.0 +_PROJECT_ATOL = 0.0 +_MAX_REFERENCE_RELATIVE_ERROR = 0.1 + + +def _offset_rms_norm_reference( + input: torch.Tensor, + weight: torch.Tensor, +) -> torch.Tensor: + input_dtype = input.dtype + input_fp32 = input.float() + inverse_rms = torch.rsqrt(input_fp32.square().mean(-1, keepdim=True) + _EPS) + return ((1.0 + weight.float()) * input_fp32 * inverse_rms).to(input_dtype) + + +def _offset_rms_norm_golden( + input: torch.Tensor, + weight: torch.Tensor, +) -> torch.Tensor: + input_fp64 = input.double() + inverse_rms = torch.rsqrt(input_fp64.square().mean(-1, keepdim=True) + _EPS) + return (1.0 + weight.double()) * input_fp64 * inverse_rms + + +def _max_abs(tensor: torch.Tensor) -> float: + return tensor.abs().max().item() if tensor.numel() else 0.0 + + +def _assert_matches_golden( + testcase: unittest.TestCase, + *, + name: str, + golden: torch.Tensor, + reference: torch.Tensor, + target: torch.Tensor, +) -> None: + testcase.assertEqual(target.shape, reference.shape, name) + testcase.assertEqual(target.dtype, reference.dtype, name) + testcase.assertTrue(torch.equal(torch.isnan(target), torch.isnan(reference)), name) + testcase.assertTrue( + torch.equal(torch.isposinf(target), torch.isposinf(reference)), name + ) + testcase.assertTrue( + torch.equal(torch.isneginf(target), torch.isneginf(reference)), name + ) + + golden_fp64 = golden.double() + reference_fp64 = reference.double() + target_fp64 = target.double() + reference_error = _max_abs(reference_fp64 - golden_fp64) + target_error = _max_abs(target_fp64 - golden_fp64) + rounding_floor = _max_abs(golden_fp64.to(target.dtype).double() - golden_fp64) + absolute_floor = max(_PROJECT_ATOL, rounding_floor) + threshold = _FUDGE_FACTOR * reference_error + absolute_floor + golden_scale = _max_abs(golden_fp64) + reference_relative_error = reference_error / max( + golden_scale, + absolute_floor, + torch.finfo(torch.float64).tiny, + ) + testcase.assertLessEqual( + reference_relative_error, + _MAX_REFERENCE_RELATIVE_ERROR, + f"{name}: reference is too inaccurate to gate the target", + ) + testcase.assertLessEqual( + target_error, + threshold, + ( + f"{name}: target_error={target_error:.6e}, " + f"reference_error={reference_error:.6e}, " + f"rounding_floor={rounding_floor:.6e}, threshold={threshold:.6e}" + ), + ) + + +@unittest.skipUnless(torch.cuda.is_available(), "CUDA required") +class TestTritonOffsetRMSNormNumerics(unittest.TestCase): + def _run_case( + self, + shape: tuple[int, ...], + dtype: torch.dtype, + seed: int, + *, + scale: float = 1.0, + ) -> None: + generator = torch.Generator(device="cuda").manual_seed(seed) + input_data = ( + torch.randn(shape, device="cuda", dtype=dtype, generator=generator) * scale + ) + weight_data = torch.randn( + shape[-1], + device="cuda", + dtype=dtype, + generator=generator, + ) + grad_output_data = torch.randn( + shape, + device="cuda", + dtype=dtype, + generator=generator, + ) + + golden_input = input_data.double().requires_grad_() + golden_weight = weight_data.double().requires_grad_() + golden_output = _offset_rms_norm_golden(golden_input, golden_weight) + golden_grads = torch.autograd.grad( + golden_output, + (golden_input, golden_weight), + grad_output_data.double(), + ) + + reference_input = input_data.detach().clone().requires_grad_() + reference_weight = weight_data.detach().clone().requires_grad_() + reference_output = _offset_rms_norm_reference( + reference_input, + reference_weight, + ) + reference_grads = torch.autograd.grad( + reference_output, + (reference_input, reference_weight), + grad_output_data, + ) + + target_input = input_data.detach().clone().requires_grad_() + target_weight = weight_data.detach().clone().requires_grad_() + target_output = triton_offset_rms_norm( + target_input, + target_weight, + _EPS, + ) + target_grads = torch.autograd.grad( + target_output, + (target_input, target_weight), + grad_output_data, + ) + + _assert_matches_golden( + self, + name="output", + golden=golden_output, + reference=reference_output, + target=target_output, + ) + for name, golden_grad, reference_grad, target_grad in zip( + ("grad_input", "grad_weight"), + golden_grads, + reference_grads, + target_grads, + ): + _assert_matches_golden( + self, + name=name, + golden=golden_grad, + reference=reference_grad, + target=target_grad, + ) + + def test_forward_and_backward_against_golden(self): + cases = ( + ((3, 255), torch.float32, 1), + ((3, 256), torch.bfloat16, 2), + ((3, 257), torch.bfloat16, 3), + ((3, 257), torch.float16, 10), + ((8, 6, 256), torch.bfloat16, 4), + ((8, 4095), torch.bfloat16, 5), + ((8, 4096), torch.bfloat16, 6), + ((8, 4097), torch.bfloat16, 7), + ((8, 5120), torch.bfloat16, 8), + ) + for shape, dtype, seed in cases: + with self.subTest(shape=shape, dtype=dtype, seed=seed): + self._run_case(shape, dtype, seed) + + def test_near_zero_variance(self): + self._run_case((8, 5120), torch.bfloat16, 9, scale=1e-5) + + def test_zero_variance(self): + self._run_case((8, 5120), torch.bfloat16, 11, scale=0.0) + + def test_custom_op_contract(self): + input = torch.randn( + 8, + 256, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + weight = torch.randn( + 256, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + torch.library.opcheck( + _triton_offset_rms_norm_op, + (input, weight, _EPS), + test_utils=( + "test_schema", + "test_faketensor", + "test_autograd_registration", + ), + ) + output, inverse_rms = _triton_offset_rms_norm_op(input, weight, _EPS) + torch.library.opcheck( + _triton_offset_rms_norm_backward_op, + (torch.randn_like(output), input, weight, inverse_rms), + test_utils=("test_schema", "test_faketensor"), + ) + + def test_torch_compile(self): + input = torch.randn( + 8, + 5120, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + weight = torch.randn( + 5120, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + grad_output = torch.randn_like(input) + compiled = torch.compile(triton_offset_rms_norm, fullgraph=True) + + expected = triton_offset_rms_norm(input, weight, _EPS) + expected_grads = torch.autograd.grad( + expected, + (input, weight), + grad_output, + ) + + compiled_input = input.detach().clone().requires_grad_() + compiled_weight = weight.detach().clone().requires_grad_() + actual = compiled(compiled_input, compiled_weight, _EPS) + actual_grads = torch.autograd.grad( + actual, + (compiled_input, compiled_weight), + grad_output, + ) + + torch.testing.assert_close(actual, expected) + for actual_grad, expected_grad in zip(actual_grads, expected_grads): + torch.testing.assert_close(actual_grad, expected_grad) + + +if __name__ == "__main__": + unittest.main() diff --git a/torchtitan/components/checkpointer/base.py b/torchtitan/components/checkpointer/base.py index 70a8344243..a9ac0a670f 100644 --- a/torchtitan/components/checkpointer/base.py +++ b/torchtitan/components/checkpointer/base.py @@ -75,15 +75,16 @@ def purge_thread( def _shares_storage(a: torch.Tensor, b: torch.Tensor) -> bool: """Whether ``a`` and ``b`` are backed by the same storage. - For ``DTensor`` the local shard's storage is compared via ``_local_tensor`` - rather than ``to_local()``, which is autograd-aware; this is a read-only - identity check on the local storage. + For ``DTensor`` the local shard is compared via ``_local_tensor`` rather + than ``to_local()``, which is autograd-aware. The dispatcher-level alias + check also supports wrapper subclasses without directly accessible storage. """ if isinstance(a, DTensor): a = a._local_tensor if isinstance(b, DTensor): b = b._local_tensor - return a.untyped_storage().data_ptr() == b.untyped_storage().data_ptr() + # pyrefly: ignore [missing-attribute] + return torch._C._is_alias_of(a, b) class ModelWrapper(Stateful): @@ -197,6 +198,9 @@ class BaseCheckpointManager(Configurable, ABC): """ enable: bool + load_only: bool + interval: int + enable_first_step_checkpoint: bool staging_future: Future | None save_future: Future | None folder: str @@ -229,7 +233,7 @@ def save(self, curr_step: int, last_step: bool = False) -> bool: def maybe_wait_for_staging(self) -> None: """Block until asynchronous staging for the last save completes.""" - if not self.enable or getattr(self, "staging_future", None) is None: + if not self.enable: return self._maybe_wait_for_staging() @@ -240,9 +244,11 @@ def close(self) -> None: # raised before assigning ``enable``. if not getattr(self, "enable", False): return - self.maybe_wait_for_staging() - self.maybe_wait_for_saving() - self._close() + try: + self.maybe_wait_for_staging() + self.maybe_wait_for_saving() + finally: + self._close() def maybe_wait_for_saving(self) -> None: """Block until the last asynchronous save completes. @@ -258,6 +264,23 @@ def maybe_wait_for_saving(self) -> None: def _wait_for_saving(self) -> None: """Await ``save_future`` and clear it. Only called when it is set.""" + # Policies shared by every manager. These depend only on config fields that + # BaseCheckpointManager.Config declares, not on how a backend reads or + # writes bytes, so they live here rather than once per backend. + + def _should_save(self, curr_step: int, last_step: bool = False) -> bool: + """Whether ``curr_step`` is a checkpointing step.""" + if not self.enable or self.load_only: + return False + if curr_step == 1 and self.enable_first_step_checkpoint: + return True + return last_step or curr_step % self.interval == 0 + + def _create_checkpoint_id(self, step: int, folder: str = "") -> str: + """Standardized checkpoint path, e.g. ``checkpoints/step-100``.""" + folder = folder or self.folder + return filesystem.join(folder, f"step-{step}") + @abstractmethod def _load(self, step: int = -1) -> bool: """Implement ``load``. Only called when checkpointing is enabled.""" @@ -328,31 +351,61 @@ def _find_load_step(self, folder: str = "") -> int: valid_steps.append(step) return max(valid_steps) if valid_steps else -1 - def _purge_stale_checkpoints(self) -> None: - """Delete the checkpoints beyond the ``keep_latest_k`` most recent.""" - if not self._should_purge(): - return + def _purge_stale_checkpoints( + self, + *, + saving_step: int, + staging_dir_prefix: str | None = None, + ) -> None: + """Delete abandoned entries and reserve one retained slot for this save.""" + if self._should_purge(): + saving_dirnames = {f"step-{saving_step}"} + if staging_dir_prefix: + saving_dirnames.add(f"{staging_dir_prefix}step-{saving_step}") + + staging_pattern = ( + re.compile(rf"{re.escape(staging_dir_prefix)}step-(0|[1-9]\d*)") + if staging_dir_prefix + else None + ) + checkpoints: list[tuple[int, str]] = [] + abandoned: list[str] = [] + + for dirname in self._storage.listdir(self.folder): + if dirname in saving_dirnames: + continue + + checkpoint_dir = filesystem.join(self.folder, dirname) + # torch_checkpointing uses this pattern for staging directories. + if staging_pattern and staging_pattern.fullmatch(dirname): + abandoned.append(checkpoint_dir) + continue + + step = self._parse_step(dirname) + if step is None: + continue + if self._is_valid_checkpoint(checkpoint_dir): + checkpoints.append((step, checkpoint_dir)) + else: + abandoned.append(checkpoint_dir) + + checkpoints.sort() + num_to_keep = self.keep_latest_k - 1 + num_to_purge = max(0, len(checkpoints) - num_to_keep) + for step, checkpoint_dir in checkpoints[:num_to_purge]: + if self._is_purge_exempt(step): + logger.info( + "Checkpointer is preserving checkpoint %s outside " + "keep_latest_k.", + checkpoint_dir, + ) + continue + assert self.purge_thread is not None + self.purge_queue.put(checkpoint_dir) - discovered: list[tuple[int, str]] = [] - for filename in self._storage.listdir(self.folder): - step = self._parse_step(filename) - if step is None: - continue - checkpoint_id = filesystem.join(self.folder, filename) - if self._is_valid_checkpoint(checkpoint_id): - discovered.append((step, checkpoint_id)) - - discovered.sort() - for step, path in discovered[: -self.keep_latest_k]: - if self._is_purge_exempt(step): - logger.info( - "Checkpointer is preserving checkpoint %s outside " - "keep_latest_k.", - path, - ) - continue - assert self.purge_thread is not None - self.purge_queue.put(path) + for checkpoint_dir in abandoned: + assert self.purge_thread is not None + self.purge_queue.put(checkpoint_dir) @dataclass(kw_only=True, slots=True) class Config(Configurable.Config): diff --git a/torchtitan/components/checkpointer/dcp.py b/torchtitan/components/checkpointer/dcp.py index 5ecd3c296c..a97061ceba 100644 --- a/torchtitan/components/checkpointer/dcp.py +++ b/torchtitan/components/checkpointer/dcp.py @@ -427,6 +427,7 @@ def _save(self, curr_step: int, last_step: bool = False) -> bool: sl.add_step_tag("checkpoint_save") self.maybe_wait_for_saving() + self._purge_stale_checkpoints(saving_step=curr_step) begin = time.monotonic() checkpoint_phase = ( @@ -487,8 +488,6 @@ def _save(self, curr_step: int, last_step: bool = False) -> bool: enable_garbage_collection=True, ) - self._purge_stale_checkpoints() - logger.info( f"Finished {checkpoint_phase} the checkpoint in " f"{time.monotonic() - begin:.2f} seconds." @@ -671,12 +670,6 @@ def _is_valid_checkpoint(self, checkpoint_dir: str) -> bool: ) ) - def _create_checkpoint_id(self, step: int, folder: str = "") -> str: - """Generate the standardized filesystem path for a checkpoint - (e.g., 'checkpoints/step-100').""" - folder = folder or self.folder - return filesystem.join(folder, f"step-{step}") - def _flattened_model_states_sd( self, state_dict: dict[str, Any] | None = None ) -> dict[str, Any]: @@ -778,21 +771,3 @@ def _save_last_step(self, curr_step: int) -> None: enable_garbage_collection=True, to_hf=self.last_save_in_hf, ) - - def _should_save(self, curr_step: int, last_step: bool = False) -> bool: - """Determine whether a checkpoint should be saved based on - the current step, interval, and training status.""" - - if not self.enable or self.load_only: - return False - - if curr_step == 1 and self.enable_first_step_checkpoint: - return True - - if last_step: - return True - - if curr_step % self.interval == 0: - return True - - return False diff --git a/torchtitan/components/checkpointer/torch_checkpointing.py b/torchtitan/components/checkpointer/torch_checkpointing.py index a559886733..e8c91323e4 100644 --- a/torchtitan/components/checkpointer/torch_checkpointing.py +++ b/torchtitan/components/checkpointer/torch_checkpointing.py @@ -7,30 +7,42 @@ from __future__ import annotations import os +import queue +import threading +from concurrent.futures import Future from dataclasses import dataclass from pathlib import Path from typing import Any +import torch import torch.nn as nn +from torch.distributed.checkpoint.state_dict_saver import _stateful_to_state_dict from torch_checkpointing.barriers import TCPStoreBarrierConfig from torch_checkpointing.checkpoint_manager import ( CheckpointManager as BackendCheckpointManager, ) from torch_checkpointing.checkpoint_writer import CheckpointWriterConfig -from torch_checkpointing.config import AsyncCheckpointSaverConfig +from torch_checkpointing.config import ( + AsyncCheckpointSaverConfig, + CheckpointSaverConfig, + SyncCheckpointSaverConfig, +) from torch_checkpointing.default_resharder import DefaultResharder from torch_checkpointing.distributed_metadata import ( METADATA_FILE_NAME as TORCH_CHECKPOINTING_METADATA_FILE_NAME, ) from torch_checkpointing.schema import ItemSpec from torch_checkpointing.staging import CheckpointStagerConfig -from torch_checkpointing.storage.base_storage import Storage +from torch_checkpointing.storage.base_storage import Storage, StorageConfig from torch_checkpointing.storage.filesystem import LocalFileSystemStorageConfig from torchtitan.components.data.loader import BaseDataLoader from torchtitan.components.optimizer import LRSchedulersContainer, OptimizersContainer from torchtitan.config import TORCH_DTYPE_MAP +from torchtitan.observability import structured_logger as sl from torchtitan.protocols.state_dict_adapter import BaseStateDictAdapter from torchtitan.tools import filesystem +from torchtitan.tools.logging import logger +from torchtitan.tools.utils import GarbageCollection from .base import ( BaseCheckpointManager, @@ -39,6 +51,7 @@ MODEL, ModelWrapper, OPTIMIZER, + purge_thread, ) DEFAULT_TORCH_CHECKPOINTING_BARRIER_TCPSTORE_PORT = 43001 @@ -97,34 +110,69 @@ def _item_specs() -> dict[str, ItemSpec]: } -def _default_backend_config() -> BackendCheckpointManager.Config: - barrier_timeout_sec = _DEFAULT_BARRIER_TIMEOUT_SEC - save_config = AsyncCheckpointSaverConfig( - writer_config=CheckpointWriterConfig( - checkpoint_write_barrier_timeout_sec=barrier_timeout_sec, - barrier_config=TCPStoreBarrierConfig( +def _writer_config(*, use_barrier: bool) -> CheckpointWriterConfig: + return CheckpointWriterConfig( + checkpoint_write_barrier_timeout_sec=_DEFAULT_BARRIER_TIMEOUT_SEC, + barrier_config=( + TCPStoreBarrierConfig( master_address=os.environ.get("MASTER_ADDR", "localhost"), tcpstore_port=DEFAULT_TORCH_CHECKPOINTING_BARRIER_TCPSTORE_PORT, timeout_barrier_init_sec=_DEFAULT_BARRIER_INIT_TIMEOUT_SEC, use_checkpoint_barrier_tcpstore_libuv=True, - ), + ) + if use_barrier + else None ), + ) + + +def _async_save_config() -> AsyncCheckpointSaverConfig: + return AsyncCheckpointSaverConfig( + writer_config=_writer_config(use_barrier=True), staging_config=CheckpointStagerConfig(use_pinned_memory=True), - wait_timeout_secs=barrier_timeout_sec, + wait_timeout_secs=_DEFAULT_BARRIER_TIMEOUT_SEC, ) + + +def _sync_save_config(*, use_barrier: bool = True) -> SyncCheckpointSaverConfig: + return SyncCheckpointSaverConfig( + writer_config=_writer_config(use_barrier=use_barrier), + wait_timeout_secs=_DEFAULT_BARRIER_TIMEOUT_SEC, + ) + + +def _default_backend_config( + save_config: CheckpointSaverConfig, + *, + storage_config: StorageConfig | None = None, +) -> BackendCheckpointManager.Config: return BackendCheckpointManager.Config( items=_item_specs(), default=ItemSpec(requires_copy=False), save=save_config, + storage_config=storage_config, ) class TorchCheckpointingManager(BaseCheckpointManager): - """TorchTitan checkpoint manager backed by ``torch_checkpointing``.""" + """TorchTitan checkpoint manager backed by ``torch_checkpointing``. + + Args: + storage_config: Backend storage for reading and writing checkpoints. + Defaults to the local filesystem. An init parameter rather than a + ``Config`` field because ``Configurable.Config`` is Tyro-parsed and + a backend storage object is not a command-line surface; callers that + need remote storage pass it programmatically. + """ @dataclass(kw_only=True, slots=True) class Config(BaseCheckpointManager.Config): - pass + def __post_init__(self) -> None: + BaseCheckpointManager.Config.__post_init__(self) + if self.last_save_in_hf: + raise ValueError( + "TorchCheckpointingManager does not support last_save_in_hf yet." + ) def __init__( self, @@ -137,11 +185,13 @@ def __init__( states: dict[str, Any], sd_adapter: BaseStateDictAdapter | None, base_folder: str = "", + storage_config: StorageConfig | None = None, ) -> None: self.enable = config.enable if not self.enable: return - self.save_future = None + self.save_future: Future[Any] | None = None + self.purge_thread: threading.Thread | None = None self.folder = filesystem.join(base_folder, config.folder) # Checked here, not just in the storage adapter: a save runs no path @@ -182,38 +232,89 @@ def __init__( self.purge_exempt = ( config.purge_exempt.build() if config.purge_exempt is not None else None ) + + save_config = ( + _sync_save_config(use_barrier=False) + if self.load_only + else _async_save_config() + ) + manager_config = _default_backend_config( + save_config, + storage_config=storage_config, + ) + self._manager_config = manager_config + storage_config = ( + self._manager_config.storage_config or LocalFileSystemStorageConfig() + ) + self._storage = _BackendCheckpointStorage(storage_config.create_storage()) + self._prewarmed = False + self.sd_adapter = sd_adapter if self.last_save_in_hf and self.sd_adapter is None: raise ValueError( "checkpoint.last_save_in_hf is True, but sd_adapter is not provided." ) - manager_config = _default_backend_config() - storage_config = manager_config.storage_config or LocalFileSystemStorageConfig() - self._storage = _BackendCheckpointStorage(storage_config.create_storage()) - self._manager = manager_config.build() + self._manager = self._manager_config.build() - def __del__(self) -> None: - self.close() + if self.keep_latest_k > 0: + self.purge_queue: queue.Queue[str | None] = queue.Queue() + self.purge_thread = threading.Thread( + target=purge_thread, + args=(self.purge_queue, self._storage.remove), + daemon=True, + ) + self.purge_thread.start() + + logger.info( + "Checkpointing active. Checkpoints will be loaded from and saved " + f"to {self.folder}" + ) - # Save and load routing land in later changes; this one only plumbs config. - # The methods are stubbed rather than omitted because BaseCheckpointManager - # declares them abstract, so a partial implementation cannot be instantiated. + def __del__(self) -> None: + # __init__ can fail before the backend manager is built. In that case, + # this object owns no backend resources to close. + if hasattr(self, "_manager"): + self.close() + # Load routing lands in a later change. def _load(self, step: int = -1) -> bool: raise NotImplementedError( "TorchCheckpointingManager does not implement load() yet." ) + @sl.log_trace_span("checkpoint_save") + @torch.no_grad() def _save(self, curr_step: int, last_step: bool = False) -> bool: - raise NotImplementedError( - "TorchCheckpointingManager does not implement save() yet." + should_save = self._should_save(curr_step, last_step) + # Prewarm on a step we are not saving, so the first real save does not + # pay for pinned-buffer allocation. + if not should_save and self._should_prewarm(): + self._manager.prewarm_staging(_stateful_to_state_dict(self.states)) + self._prewarmed = True + if not should_save: + return False + + sl.add_step_tag("checkpoint_save") + self.maybe_wait_for_saving() + # Always preserve the current step's published and staging directories. + self._purge_stale_checkpoints( + saving_step=curr_step, + staging_dir_prefix=( + self._manager_config.save.writer_config.temp_dir_prefix + ), ) - def _wait_for_saving(self) -> None: - raise NotImplementedError( - "TorchCheckpointingManager does not implement saving yet." - ) + if last_step: + self._save_last_step(curr_step) + else: + self.save_future = self._manager.save( + self._create_checkpoint_id(curr_step), + _stateful_to_state_dict(self.states), + ) + self._prewarmed = True + + return True def _is_valid_checkpoint(self, checkpoint_dir: str) -> bool: return self._storage.isfile( @@ -221,13 +322,69 @@ def _is_valid_checkpoint(self, checkpoint_dir: str) -> bool: ) def _maybe_wait_for_staging(self) -> None: - raise NotImplementedError( - "TorchCheckpointingManager does not implement maybe_wait_for_staging() yet." - ) + # BaseCheckpointManager.close() calls this to wait for in-flight staging. + # If _save_last_step already closed the backend manager, that close drained + # staging but left this lock usable, so acquiring it here cannot hang. + with self._manager.lock(): + pass + + def _wait_for_saving(self) -> None: + # Clear the active save before waiting so close() does not retry a failure. + save_future = self.save_future + assert save_future is not None + self.save_future = None + save_future.result(timeout=self._manager_config.save.wait_timeout_secs) def _close(self) -> None: - # hasattr: __del__ -> close() can reach here on a partially constructed - # object if __init__ raised after setting enable but before building the - # backend manager. - if hasattr(self, "_manager"): - self._manager.close() + try: + self.maybe_wait_for_saving() + finally: + try: + if self.purge_thread is not None and self.purge_thread.is_alive(): + self.purge_queue.put(None) + self.purge_thread.join() + finally: + # _save_last_step may already have closed the manager; the + # backend's close() returns immediately when it has. + self._manager.close() + + def _save_last_step(self, curr_step: int) -> None: + if self.last_save_model_only: + model_state = self.states[MODEL].state_dict() + # Cast floating-point tensors to the export dtype and preserve other + # buffers. + model_state = { + key: value.to(self.export_dtype) + if isinstance(value, torch.Tensor) + and value.is_floating_point() + and value.dtype != self.export_dtype + else value + for key, value in model_state.items() + } + states: dict[str, Any] = {MODEL: model_state} + logger.info( + f"Saving a model only checkpoint in {self.export_dtype} " + f"at last step, step {curr_step}." + ) + else: + states = self.states + logger.info(f"Saving a full checkpoint at last step, step {curr_step}.") + + # The final save must land before the process exits, so retire the async + # manager and write synchronously through a fresh one. + self._manager.close() + manager = _default_backend_config( + _sync_save_config(), + storage_config=self._manager_config.storage_config, + ).build() + try: + manager.save( + self._create_checkpoint_id(curr_step), + _stateful_to_state_dict(states), + ) + finally: + manager.close() + GarbageCollection.collect("GC collection invoked by checkpointer.") + + def _should_prewarm(self) -> bool: + return self.enable and not self._prewarmed and not self.load_only diff --git a/torchtitan/components/quantization/__init__.py b/torchtitan/components/quantization/__init__.py index 91986ec3c5..0537502c04 100644 --- a/torchtitan/components/quantization/__init__.py +++ b/torchtitan/components/quantization/__init__.py @@ -37,7 +37,7 @@ class Config(ModelConfigConverter.Config): Float8Linear, Float8LinearConverter, ) -from .mx import ( # noqa: F401, E402 +from .mxfp8 import ( # noqa: F401, E402 MXFP8GroupedExpertsConverter, MXFP8Linear, MXFP8LinearConverter, diff --git a/torchtitan/components/quantization/_fsdp_tensor.py b/torchtitan/components/quantization/_fsdp_tensor.py new file mode 100644 index 0000000000..76dd8b683c --- /dev/null +++ b/torchtitan/components/quantization/_fsdp_tensor.py @@ -0,0 +1,569 @@ +# 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. + +"""Private FSDP shard/unshard lifecycle for quantized parameters. + +The lifecycle uses two tensor subclasses, one per state: + +``_ShardedFSDPTensor`` + The persistent parameter. Holds the high-precision shard, owns the FSDP + pre/post-all-gather hooks, and knows how to build a format's unsharded + operands. A format subclasses this and implements one method. + +``_UnshardedFSDPTensor`` + The unsharded tensor for one unshard lifetime. Holds only the format's + operands and no high-precision storage, which is what lets FSDP release + the all-gather output. Generic: it derives its unsharded inner tensors from the + operands dataclass fields, so formats do not subclass it. + +Both present the logical high-precision metadata of the model parameter, so +autograd returns a high-precision parameter gradient in either state. + +Instance timeline, from module construction to reshard:: + + module __init__ S = _ShardedFSDPTensor(storage_shard) + | one instance, lives for the whole run; + | it is the nn.Parameter and the checkpoint + v + ---- unshard ---------------------------------------------------------- + fsdp_pre_all_gather S._tensor.to(param_dtype) -> comm tensor + | + v + (all-gather) replicated param_dtype tensor, temporary + | + fsdp_post_all_gather | out is None: first unshard + S builds -------> C = _UnshardedFSDPTensor(operands) + | new instance; holds only qdata/scales. + | Returned to FSDP with the operands' + | tensors so FSDP can manage their storage. + v + comm tensor released -- C never referenced it + | + forward/backward compute reads C.operands + | + ---- reshard ---------------------------------------------------------- + FSDP frees the storage of C's unsharded inner tensors. + C itself stays alive: autograd and the module may + still hold it, and its addresses must not move. + | + ---- unshard again ---------------------------------------------------> + fsdp_post_all_gather | out is C: refill + S refills ------> C's existing tensors are written in place + | no new instance; _validate_refilled_tensor_ + | identity() enforces that + v + (repeats until the final reshard) + +So S is created once per parameter and C once per *distinct* unshard +lifetime -- not once per unshard. RAF=False keeps a single C alive across +forward, backward, recomputation, and pipeline microbatches; RAF=True +reuses that same C object, refilling its storage before backward. + +GraphTrainer's SimpleFSDP reaches the same place by a different route: it +reconstructs the replicated tensor itself, then quantizes and wraps it from +inside its own parametrization, constructing C directly. That path owns the +gradient edge FSDP2 creates internally, so it lives with SimpleFSDP rather +than here. + +Two dtypes appear above and they are not the same one. S holds the parameter's +own storage dtype, set by ``training.dtype`` -- float32 by default, so S is +normally an fp32 master weight, not a BF16 one. ``mp_policy.param_dtype`` +(``training.mixed_precision_param``, bfloat16 by default) is only what +``fsdp_pre_all_gather`` casts *to*, so it is the dtype of the comm tensor, the +all-gather output, and hence the logical tensor a format quantizes. Nothing +here requires either to be BF16; MXFP8 separately rejects a non-BF16 weight in +``_MXFP8LinearFunction.forward``, because its kernels need one. + +Terminology +----------- + +Several names here are a word apart, so: + +operands + A format's quantized output: a frozen dataclass of qdata, scales, any + workspace. Plain data; not a tensor subclass. + +unsharded tensor + The ``_UnshardedFSDPTensor`` that *holds* the operands and presents the + parameter's logical high-precision metadata. "C" below. + +``_build_operands(logical_tensor, out=None)`` + The format's quantizer, and the only method a format implements. Returns + a fresh operands dataclass, or with ``out`` set refills that one's + existing tensors in place instead of allocating. + +logical tensor + The unsharded high-precision tensor handed to the quantizer. "Logical" + because it matches the size the model declares, excluding any padding the + all-gather added to make every shard the same size. + +unsharded inner tensors + The operands' dataclass *fields*, whose storage FSDP allocates, frees, and + refills across the unshard lifecycle. A strict subset of what the operands + expose: a derived view belongs in a property, and a property is never one + of these, since FSDP must not free the same allocation twice. They are + exactly the tensors ``__tensor_flatten__`` reports, which FSDP holds as + ``FSDPParam._unsharded_inner_tensors`` and frees in + ``free_unsharded_param``. + +metadata source + The tensor an unsharded tensor copies its shape, dtype, device, and + layout from -- normally the unsharded high-precision tensor it was built + from, since an unsharded tensor has no storage of its own to describe. + +Adding a format +--------------- + +Subclass ``_ShardedFSDPTensor``; do not subclass ``_UnshardedFSDPTensor``. + +A format differs from every other format in exactly one way: how it turns a +high-precision tensor into its operands. That belongs on the sharded class, +because the sharded tensor is the parameter FSDP calls hooks on, so it is +what *produces* the operands. ``_UnshardedFSDPTensor`` only *holds* them, and +holding is format-independent -- it reads the operand tensors off the +operands dataclass fields, which works for any format. Subclassing +it would add a type that overrides nothing. + +So a format supplies a frozen dataclass of the tensors one unshard lifetime +owns, and one method, ``_build_operands``. Everything else -- +flattening, refill, reshard, the SimpleFSDP bridge -- comes from here. +""" + +from __future__ import annotations + +import math +from dataclasses import fields, is_dataclass +from typing import Any + +import torch +from torch.utils import _pytree as pytree +from torch.utils._python_dispatch import return_and_correct_aliasing + + +# The unsharded-tensor machinery is internal. A data-parallel implementation +# reaches it through the FSDP hooks, or by calling +# ``_build_operands`` itself when it reconstructs the +# unsharded tensor, as GraphTrainer's SimpleFSDP does. +__all__: list[str] = [] + +# Ops FSDP performs on the sharded parameter's real storage. The wrapper must +# survive them so the parameter keeps its identity across FSDP bookkeeping. +_FSDP_SHARDED_OPS = { + torch.ops.aten.empty_like.default, + torch.ops.aten.new_zeros.default, + torch.ops.aten.slice.Tensor, + torch.ops.aten.copy_.default, + torch.ops.aten.view.default, + torch.ops.aten.as_strided.default, + torch.ops.aten._to_copy.default, + torch.ops.aten._pin_memory.default, + torch.ops.aten.split.Tensor, + torch.ops.aten.clone.default, + torch.ops.aten.transpose.int, + torch.ops.aten.t.default, + torch.ops.c10d.scatter_.default, + torch.ops.aten.detach.default, + torch.ops.aten.alias.default, +} + +# The unsharded tensor has no high-precision storage, so only metadata-level ops +# are answerable. Views re-wrap; factories allocate fresh plain tensors. +_FSDP_UNSHARDED_VIEW_OPS = { + torch.ops.aten.alias.default, + torch.ops.aten.as_strided.default, + torch.ops.aten.detach.default, + torch.ops.aten.view.default, +} + +_FSDP_UNSHARDED_FACTORY_OPS = { + torch.ops.aten.empty_like.default, + torch.ops.aten.new_zeros.default, + torch.ops.aten.zeros_like.default, +} + + +def _unsharded_inner_tensor_names(operands_cls: type) -> tuple[str, ...]: + """Return the names of the unsharded inner tensors an operands dataclass owns.""" + if not is_dataclass(operands_cls): + raise TypeError( + "An operands must be a dataclass of tensors; got " + f"{operands_cls.__name__}." + ) + return tuple(field.name for field in fields(operands_cls)) + + +def _unsharded_inner_tensors(operands: Any) -> tuple[torch.Tensor, ...]: + """Return the unsharded inner tensors an operands dataclass owns. + + Field order defines the unsharded inner tensor order, so no format restates + it. Every field must be a distinct allocation: FSDP takes ownership of + each one's storage, so listing two views of the same storage would make it + free the same memory twice and leave ``__tensor_unflatten__`` unable to + tell which field was the derived one. Derived views belong in properties, + which ``fields()`` skips. + """ + tensors = tuple( + getattr(operands, name) + for name in _unsharded_inner_tensor_names(type(operands)) + ) + storages = {tensor.untyped_storage()._cdata for tensor in tensors} + if len(storages) != len(tensors): + raise ValueError( + f"{type(operands).__name__} fields must be distinct " + "allocations; a field aliasing another field's storage should be " + "a property instead." + ) + return tensors + + +def _validate_refilled_tensor_identity( + unsharded_inner_tensors: tuple[torch.Tensor, ...], + refilled_tensors: tuple[torch.Tensor, ...], +) -> None: + """Require a refill to preserve every tensor object managed by FSDP.""" + if len(unsharded_inner_tensors) != len(refilled_tensors) or any( + previous is not current + for previous, current in zip( + unsharded_inner_tensors, refilled_tensors, strict=True + ) + ): + raise RuntimeError( + "FSDP unsharded-operands refill replaced inner tensor storage" + ) + + +class _FSDPTensorBase(torch.Tensor): + """Logical high-precision metadata shared by both lifecycle states.""" + + @staticmethod + def __new__(cls, tensor: torch.Tensor, *args: Any, **kwargs: Any): + del args + return torch.Tensor._make_wrapper_subclass( + cls, + kwargs.get("_logical_size", tensor.size()), + strides=kwargs.get("_logical_stride", tensor.stride()), + storage_offset=kwargs.get( + "_logical_storage_offset", tensor.storage_offset() + ), + dtype=kwargs.get("_logical_dtype", tensor.dtype), + layout=tensor.layout, + device=kwargs.get("_logical_device", tensor.device), + pin_memory=tensor.is_pinned(), + requires_grad=kwargs.get("_logical_requires_grad", tensor.requires_grad), + ) + + +class _ShardedFSDPTensor(_FSDPTensorBase): + """Persistent high-precision parameter that owns the FSDP hooks. + + This is the sharded half of the lifecycle. It is the ``nn.Parameter`` the + optimizer updates and the checkpoint stores, it holds the high-precision + shard in ``_tensor``, and it lives for the whole run. It is *not* what + compute sees under FSDP: the post-all-gather hook hands back a + :class:`_UnshardedFSDPTensor` holding the quantized operands, and that is + what forward and backward read for the duration of one unshard. See the + instance timeline at the top of this module for how the two hand off. + + **This is the class a format subclasses**, because a format differs only + in how it turns a high-precision tensor into operands, and this is the + side that produces them. ``_UnshardedFSDPTensor`` only holds them, which is + format-independent, so it is generic and is never subclassed. + + A subclass supplies: + + * a frozen dataclass of the tensors one unshard lifetime owns, whose + fields are distinct allocations -- derived views belong in properties; + * ``_build_operands(logical_tensor, out=None)``, which + allocates a new operands when ``out`` is None and otherwise + refills ``out``'s existing tensors in place. + + Everything else -- flattening, refill, reshard, and the SimpleFSDP bridge + -- is inherited. + """ + + def __init__(self, tensor: torch.Tensor, **logical_metadata: Any) -> None: + del logical_metadata + self._tensor = tensor + + def __tensor_flatten__(self): + return ["_tensor"], (self.dtype,) + + @classmethod + def __tensor_unflatten__(cls, inner_tensors, metadata, outer_size, outer_stride): + del metadata, outer_size, outer_stride + return cls(inner_tensors["_tensor"]) + + @classmethod + # pyrefly: ignore [bad-param-name-override] + def __torch_dispatch__(cls, func, types, args, kwargs=None): + del types + template = None + preserve_wrapper = func in _FSDP_SHARDED_OPS + + def unwrap(tensor: _ShardedFSDPTensor) -> torch.Tensor: + nonlocal template + if template is None: + template = tensor + elif preserve_wrapper and type(tensor) is not type(template): + raise RuntimeError("FSDP operation mixed sharded tensor types") + return tensor._tensor + + output = func( + *pytree.tree_map_only(cls, unwrap, args or ()), + **pytree.tree_map_only(cls, unwrap, kwargs or {}), + ) + if not preserve_wrapper: + return output + assert template is not None + return pytree.tree_map_only(torch.Tensor, type(template), output) + + def _build_operands( + self, + logical_tensor: torch.Tensor, + out: Any = None, + ) -> Any: + """Quantize ``logical_tensor``, into ``out``'s tensors when refilling.""" + raise NotImplementedError + + def fsdp_should_release_all_gather_outputs_after_post_all_gather(self) -> bool: + """Release the high-precision all-gather output after state construction.""" + return True + + def fsdp_pre_all_gather(self, mesh, outer_size, outer_stride, module, mp_policy): + """Return the high-precision communication tensor and the logical size. + + All-gather needs every rank to contribute the same number of elements, + so an expert count that does not divide the mesh size leaves the last + rank short. FSDP's contract is that this returns the *padded* shard; + the logical size travels in the metadata so ``fsdp_post_all_gather`` + can drop the padding before quantizing the logical tensor. + """ + del outer_stride, module + # FSDP hands the hook no shard dimension, but the local shard differs + # from the logical size exactly along it. The default, non-extension + # path has it directly as ``fsdp_placement.dim`` and rejects the same + # case there: + # https://github.com/pytorch/pytorch/blob/c7da99c173f2b67905ee798576a644b6b32cbfee/torch/distributed/fsdp/_fully_shard/_fsdp_param.py#L323-L331 + sharded_dims = [ + dim + for dim, (local, logical) in enumerate( + zip(self._tensor.shape, outer_size, strict=True) + ) + if local != logical + ] + if len(sharded_dims) > 1: + raise RuntimeError( + f"FSDP sharded more than one dimension: local " + f"{tuple(self._tensor.shape)} against logical {tuple(outer_size)}" + ) + if sharded_dims and sharded_dims[0] != 0: + raise NotImplementedError( + "FSDP unsharded tensors support sharding dimension 0 only, but " + f"this parameter of shape {tuple(outer_size)} is sharded on " + f"dimension {sharded_dims[0]}. TorchTitan selects Shard(1) for " + "grouped experts when the FSDP degree exceeds the expert " + "count, so either lower the degree or raise the expert count." + ) + dtype = mp_policy.param_dtype or self._tensor.dtype + # Pad to what FSDP calls ``padded_sharded_param_size``. The default + # path pre-pads to ``chunks[0].size()``, and torch.chunk puts the + # remainder in the earlier chunks, so that equals ceil(dim0 / world): + # https://github.com/pytorch/pytorch/blob/c7da99c173f2b67905ee798576a644b6b32cbfee/torch/distributed/fsdp/_fully_shard/_fsdp_param.py#L332-L345 + # An extension must return exactly that size; only the short ranks + # would trip the check, so the rest hang in the all-gather instead: + # https://github.com/pytorch/pytorch/blob/c7da99c173f2b67905ee798576a644b6b32cbfee/torch/distributed/fsdp/_fully_shard/_fsdp_param.py#L1143-L1158 + # + # Note the padding happens per unshard here, where the default path + # pads once. It pre-pads the sharded parameter at init and keeps that + # buffer for the run ("Pre-pad the sharded parameter to avoid padding + # before all-gather"), but an extension is handed the unpadded shard + # every time. TODO(anijain2305): hold a persistent padded buffer on the + # sharded tensor and copy into it, to drop the per-unshard allocation. + padded_rows = math.ceil(outer_size[0] / mesh.size()) + if self._tensor.size(0) != padded_rows: + # Allocate the padded buffer directly in the comm dtype and let the + # copy do the cast, rather than casting the whole shard first and + # then copying that into a second buffer. + source = self._tensor.new_zeros( + (padded_rows, *self._tensor.shape[1:]), dtype=dtype + ) + source.narrow(0, 0, self._tensor.size(0)).copy_(self._tensor) + else: + source = self._tensor.to(dtype) + return (source,), outer_size + + def fsdp_post_all_gather( + self, all_gather_outputs, metadata, param_dtype, *, out=None + ): + """Create or refill the unsharded tensor operands after all-gather.""" + del param_dtype + (logical_tensor,) = all_gather_outputs + # ``metadata`` is the logical size returned by fsdp_pre_all_gather. An + # unevenly sharded parameter gathers padding rows past it, which must + # not reach the quantizer: they would occupy real scale tiles and, for + # a grouped expert tensor, appear as extra experts. + if metadata is not None and logical_tensor.size(0) != metadata[0]: + logical_tensor = logical_tensor.narrow(0, 0, metadata[0]) + + # On the first unshard, FSDP has no unsharded-tensor container or managed + # tensors yet. Build both and return them to FSDP. With RAF=False, FSDP + # keeps these operands alive through forward and backward. + if out is None: + with torch.no_grad(): + operands = self._build_operands(logical_tensor) + return ( + _UnshardedFSDPTensor(logical_tensor, operands), + _unsharded_inner_tensors(operands), + ) + + # After FSDP releases and later unshards the tensor again, ``out`` is + # the same unsharded-tensor object returned above. This occurs between + # forward and backward with RAF=True, or after a later reshard. Refill + # the same unsharded inner tensor objects so existing module and autograd + # references remain valid. ``out`` is what this hook returned, so it is + # always the bare unsharded tensor -- FSDP2 never re-wraps it. + target = out + if not isinstance(target, _UnshardedFSDPTensor): + raise RuntimeError("FSDP output does not own operands") + existing = target.operands + unsharded_inner_tensors = _unsharded_inner_tensors(existing) + with ( + torch.no_grad(), + # Refilling lifecycle-managed storage is not a user-visible tensor + # mutation and must not invalidate saved-tensor version checks. + torch.autograd._unsafe_preserve_version_counter(unsharded_inner_tensors), + ): + refilled = self._build_operands(logical_tensor, out=existing) + _validate_refilled_tensor_identity( + unsharded_inner_tensors, _unsharded_inner_tensors(refilled) + ) + target._operands = refilled + + +class _UnshardedFSDPTensor(_FSDPTensorBase): + """Unsharded tensor holding one unshard lifetime's format operands. + + The unsharded half of the lifecycle, built by + :class:`_ShardedFSDPTensor`'s post-all-gather hook and alive until the + final reshard. Carries no high-precision storage -- that is what lets FSDP + release the all-gather output -- so reading it as a high-precision tensor + is an error; only format-aware consumers may read + ``operands``. See the instance timeline at the top of this + module for how the two classes hand off. + + Generic by design: the unsharded inner tensors come from the operands' + dataclass fields, which works for any format. Do not subclass it; formats + subclass :class:`_ShardedFSDPTensor` instead. + """ + + def __init__( + self, + metadata_source: torch.Tensor, + operands: Any, + **logical_metadata: Any, + ) -> None: + # ``__new__`` already consumed both to build the wrapper subclass with + # this tensor's logical size, stride, dtype, device and requires_grad. + # Python hands ``__init__`` the same arguments, so drop them here + # rather than store them: the metadata lives on the tensor itself. + del metadata_source, logical_metadata + self._operands = operands + # __tensor_flatten__ reports the unsharded inner tensors by attribute name and + # the subclass machinery fetches them with a plain getattr, so each has + # to exist as an attribute here -- reaching into the operands is + # not an option. Mirror rather than copy: these are the same tensor + # objects, so an in-place refill updates both views, and a refill that + # substituted objects is rejected by + # _validate_refilled_tensor_identity. + for name in _unsharded_inner_tensor_names(type(operands)): + setattr(self, f"_{name}", getattr(operands, name)) + + def __tensor_flatten__(self): + operands_cls = type(self._operands) + names = [f"_{name}" for name in _unsharded_inner_tensor_names(operands_cls)] + return names, (operands_cls, self.dtype) + + @staticmethod + def __tensor_unflatten__(inner_tensors, metadata, outer_size, outer_stride): + operands_cls, dtype = metadata + unsharded_inner_tensors = [ + inner_tensors[f"_{name}"] + for name in _unsharded_inner_tensor_names(operands_cls) + ] + operands = operands_cls(*unsharded_inner_tensors) + # FSDP supplies the logical shape; any unsharded inner tensor can stand in for + # the rest, since they share the unsharded tensor's device and layout. + return _UnshardedFSDPTensor( + unsharded_inner_tensors[0], + operands, + _logical_size=outer_size, + _logical_stride=outer_stride, + _logical_dtype=dtype, + ) + + @classmethod + # pyrefly: ignore [bad-param-name-override] + def __torch_dispatch__(cls, func, types, args, kwargs=None): + del types + template = None + + def unwrap(tensor: _UnshardedFSDPTensor) -> torch.Tensor: + nonlocal template + if template is None: + template = tensor + elif tensor._operands is not template._operands: + raise RuntimeError("FSDP operation mixed unsharded tensor operands") + # There is no high-precision storage to hand the op; a meta tensor + # carries the logical metadata that view ops need. + return torch.empty_strided( + tensor.size(), + tensor.stride(), + dtype=tensor.dtype, + device="meta", + requires_grad=tensor.requires_grad, + ) + + def wrap_view(tensor: torch.Tensor): + assert template is not None + operands = template._operands + # __new__ reads layout and pinning off a real tensor, and the + # template has no storage to answer with, so borrow a managed + # tensor for those two and give the view's logical metadata for + # everything else. Which one does not matter: they + # share the unsharded tensor's device, layout, and pinning. + layout_source = _unsharded_inner_tensors(operands)[0] + return _UnshardedFSDPTensor( + layout_source, + operands, + _logical_size=tensor.size(), + _logical_stride=tensor.stride(), + _logical_storage_offset=tensor.storage_offset(), + _logical_dtype=template.dtype, + _logical_device=template.device, + _logical_requires_grad=tensor.requires_grad, + ) + + original_args, original_kwargs = args, kwargs or {} + args, kwargs = pytree.tree_map_only( + cls, unwrap, (original_args, original_kwargs) + ) + assert template is not None + if func in _FSDP_UNSHARDED_FACTORY_OPS: + kwargs["device"] = template.device + return func(*args, **kwargs) + if func not in _FSDP_UNSHARDED_VIEW_OPS: + raise RuntimeError( + f"{func} attempted to read a storage-free FSDP unsharded tensor" + ) + wrapped = pytree.tree_map_only(torch.Tensor, wrap_view, func(*args, **kwargs)) + return return_and_correct_aliasing( + func, original_args, original_kwargs, wrapped + ) + + @property + def operands(self) -> Any: + """Return the operands for the current unshard lifetime.""" + return self._operands diff --git a/torchtitan/components/quantization/mxfp8.md b/torchtitan/components/quantization/mxfp8/README.md similarity index 50% rename from torchtitan/components/quantization/mxfp8.md rename to torchtitan/components/quantization/mxfp8/README.md index ad82edcb78..29cc0360e6 100644 --- a/torchtitan/components/quantization/mxfp8.md +++ b/torchtitan/components/quantization/mxfp8/README.md @@ -8,7 +8,10 @@ MXFP8 training can provide substantial training speedups for models where the ma - [Requirements](#requirements) - [How MXFP8 Works](#how-mxfp8-works) + - [TorchAO and TorchTitan Responsibilities](#torchao-and-torchtitan-responsibilities) + - [FSDP-Managed Dense Weights](#fsdp-managed-dense-weights) - [MXFP8 for Linear Modules](#mxfp8-for-linear-modules) + - [Input Activation Storage](#input-activation-storage) - [Usage](#usage) - [MXFP8 for Grouped GEMMs (MoE)](#mxfp8-for-grouped-gemms-moe) - [Usage](#usage-1) @@ -24,20 +27,137 @@ MXFP8 training can provide substantial training speedups for models where the ma - NVIDIA B200 (SM100 or SM100a) - PyTorch nightly -- TorchAO v0.14.0 or newer ([TorchAO Installation Guide](https://github.com/pytorch/ao#installation)) - -Note: GB200 is also supported but requires building torchao from source (see installation guide above). +- TorchAO 0.18.0 or later, with `nvidia-cutlass-dsl` and `apache-tvm-ffi` ### How MXFP8 Works MXFP8 differs from standard Float8 training in its scaling approach: - **Granular scaling factor**: Instead of using a single scale factor per tensor (tensorwise) or per row/column (rowwise), MXFP8 uses a more granular, block-based scaling with a default block size of 1x32 elements. Each block of 32 elements shares a common scale factor. The data dtype is `torch.float8_e4m3fn`, and the scale factor dtype is `torch.float8_e8mfnu`. -- **Native hardware support**: On NVIDIA B200 (Blackwell) GPUs, MXFP8 GEMMs and Grouped GEMMs are accelerated using cuBLAS and CUTLASS kernels exposed via `torch._scaled_mm` and `torch._scaled_grouped_mm`, achieving up to 2x speedup over bfloat16 on common shapes. -- **Dynamic quantization**: For every MXFP8 Linear or Grouped GEMM, activations and weights are dynamically quantized to MXFP8, then a MXFP8 GEMM/Grouped GEMM is performed, resulting in a net speedup. +- **Native hardware support**: On NVIDIA B200 (Blackwell) GPUs, MXFP8 GEMMs and Grouped GEMMs are accelerated using cuBLAS and CUTLASS kernels exposed via `torch.nn.functional.scaled_mm` and `torch._scaled_grouped_mm`, achieving up to 2x speedup over bfloat16 on common shapes. +- **Dynamic activation quantization**: Linear and Grouped GEMM activations use + standard 1x32 MXFP8 scaling and are dynamically quantized for each operation. + Forward scales are computed independently within each token row, so scale + calculation cannot carry information between causal positions. Linear WGRAD + can either retain the high-precision input and quantize it columnwise in + backward, or retain a columnwise MXFP8 operands produced in forward. + Neither choice affects causal forward outputs. +- **FSDP-managed dense weights**: After FSDP all-gathers a dense Linear weight + in BF16, TorchTitan's post-all-gather hook quantizes it with square 32x32 + scale tiles to create FPROP and DGRAD operands. FSDP owns those + buffers, so their lifetime follows the normal reshard-after-forward and + reshard-after-backward policies. The temporary BF16 all-gather output is + released after the independent MXFP8 operands is constructed. + +Dense MXFP8 linear layers combine 32x32 square scale tiles for weights with +1x32 tiles for activations. Square weight tiles still introduce quantization +error, but they are orientation-symmetric: FPROP and DGRAD use the same +quantized values and share one cached qdata allocation. This avoids choosing +two independently quantized weight operands for the two GEMM +orientations. + +#### TorchAO and TorchTitan Responsibilities + +The dense linear integration keeps a narrow boundary between TorchAO and +TorchTitan. TorchTitan uses these kernel-level operations from TorchAO: + +- `mxfp8_quantize_cuda` for rowwise and columnwise activation quantization. +- `triton_to_mxfp8_32x32_swizzle_dim0_qdata_dim01_scale` for 32x32 weight + quantization with one shared qdata allocation and both scale layouts. +- `triton_mx_block_rearrange` for scale layout conversion. + +TorchTitan owns the pieces coupled to the training system: + +- MXFP8 linear autograd. +- The generic FSDP unsharded-tensor lifecycle and its MXFP8 specialization. +- Quantized-weight storage and lifetime. +- Model-specific input-activation storage policy. + +This keeps FSDP and parallelism policy in TorchTitan while allowing additional +kernel fusion to be implemented independently in TorchAO. + +#### FSDP-Managed Dense Weights + +The FSDP post-all-gather hook quantizes each unsharded BF16 weight and returns +independent MXFP8 qdata and scale tensors for FSDP to manage. There is no +separate module-level or autograd-level weight cache. The unsharded tensor +therefore follows the normal FSDP lifecycle: + +| `reshard_after_forward` | Behavior | +| --- | --- | +| `False` | Quantize on the first unshard and reuse the MXFP8 operands until FSDP releases them after backward. Pipeline parallelism can reuse them across microbatches. | +| `True` | Reshard after forward. Backward performs another BF16 all-gather and post-all-gather quantization. | + +FSDP keeps the storage-free logical tensor stable and allocates, releases, or +refills its inner qdata and scale tensors. The module keeps an ordinary +`nn.Parameter` when FSDP is not used and quantizes it dynamically. FSDP setup +installs the unsharded-tensor wrapper immediately before sharding. This reuses +the existing FSDP state machine instead of introducing another cache lifecycle +in `MXFP8Linear`. + +The hook currently constructs both FPROP and DGRAD weight operands on every +actual unshard. This is already optimal when `reshard_after_forward=False` +because the operands are constructed once and reused. With +`reshard_after_forward=True`, phase-specific construction would require FSDP to +expose whether an unshard is serving forward, backward, or checkpoint +recomputation. ### MXFP8 for Linear Modules +Dense weights always use square 32x32 scale tiles, and both local weight +dimensions must be divisible by 32. Activations use standard 1D scaling. + +#### Input Activation Storage + +Weight-gradient computation needs the linear input during backward. +`MXFP8Linear` supports two ways to retain it: + +| Save format | Forward quantization | Saved for WGRAD | Backward work | +| --- | --- | --- | --- | +| `bf16` (default) | Rowwise only | Original BF16 input | Quantize the input columnwise | +| `mxfp8` | Rowwise and columnwise | Columnwise qdata and scales | Reuse the saved MXFP8 operand | + +Saving MXFP8 reduces activation storage and avoids a backward quantization pass +only when no other operation retains the same BF16 input. If a preceding +operation already saves its output for backward, as flash attention does, the +BF16 tensor is already available as the linear input. Saving an additional +MXFP8 operands would then increase peak memory. Since this ownership is +model-dependent, `bf16` is the conservative default and audited modules opt in +through `linears_saving_inputs_for_backward_in_mxfp8`. + +##### Interaction with activation checkpointing + +Without activation checkpointing, the selected operands remains live +from forward until WGRAD. Saving BF16 is preferable when another operation +already retains the same tensor; otherwise, saving MXFP8 can replace that BF16 +storage and avoid columnwise quantization during backward. + +With full activation checkpointing, neither operands is retained from +the original forward; it is reconstructed during backward recomputation. The +current implementation applies the configured format to both executions. An +MXFP8-selected linear therefore produces rowwise and columnwise operands +during the original forward, discards the columnwise result, and produces it +again during recomputation. Ideally, it would produce only the rowwise operand +in the original forward and produce both operands during recomputation. We keep +one policy for both executions to avoid adding recomputation detection and a +second execution-dependent autograd contract. + +More granular `torch.remat` policies introduce additional choices about which +operations and tensors are saved or recomputed. The optimal format is therefore +both model-dependent and activation-checkpointing-policy-dependent. The current +model policy deliberately remains fixed across those modes. + +The built-in policies currently select: + +| Model | Modules saving MXFP8 inputs | +| --- | --- | +| Llama 3 | `attention.qkv_linear.wqkv`, `feed_forward.w2` | +| DeepSeek V3 | `attention.wkv_b`, `feed_forward.w2`, `shared_experts.w2` | +| Flux | None | + +All other converted linears save BF16 inputs. Trainer and GraphTrainer share +these model policies. + #### Usage Quantization is applied at config time in your `model_registry()` function via the `quantization` parameter. Each converter walks the model config tree and swaps config types so that quantized modules are built directly. @@ -52,23 +172,15 @@ model_spec = model_registry( "flux-schnell", quantization=[ MXFP8LinearConverter.Config( - recipe_name="mxfp8_rceil", fqns=["double_blocks", "single_blocks"], + # Add audited single-consumer inputs here. Flux uses BF16 by default. + linears_saving_inputs_for_backward_in_mxfp8=[], model_compile_enabled=True, ), ], ) ``` -**Configuration Options:** - -* `recipe_name`: MXFP8 recipe name. Options: - * `"mxfp8_rceil"` (default): MXFP8 dynamic quantization with RCEIL rounding mode when computing the e8m0 scale factors. - * `"mxfp8_cublas"`: Use the cuBLAS-based MXFP8 recipe for best performance on B200 GPUs. - * `"mxfp8_cublas_rceil"`: Uses round-ceiling mode for scale calculation. -* `fqns` (optional): List of fully qualified names to filter which Linear modules to convert. Only `Linear.Config` entries whose FQN contains a match are swapped to `MXFP8Linear.Config`. If empty, all Linear modules are converted. -* `model_compile_enabled`: set to `True` when `torch.compile` is enabled for the model (required for competitive performance). - **Hardware Requirements:** MXFP8 training requires NVIDIA B200 (SM100) or newer GPUs. @@ -101,7 +213,6 @@ model_spec = model_registry( quantization=[ MXFP8LinearConverter.Config( - recipe_name="mxfp8_rceil", fqns=["double_blocks", "single_blocks"], model_compile_enabled=True, ), @@ -135,7 +246,6 @@ model_spec = model_registry( "671B", quantization=[ MXFP8LinearConverter.Config( - recipe_name="mxfp8_rceil", fqns=["double_blocks", "single_blocks"], model_compile_enabled=True, ), @@ -178,11 +288,7 @@ Single-node training on 8x power limited B200 GPUs, batch size 1, sequence lengt | None (bfloat16) | 6169 | - | | mxfp8 | 7401 | +20.3% | -Training runs on 64 node GB200 cluster with TorchTitan Llama4 Scout show that MXFP8 MoE training has equivalent convergence to bfloat16 training baseline. In fact, after 3,000 steps it finishes with slightly *lower* loss than bfloat16! This is consistent with our scaling experiments with [MXFP8 training for dense models](https://pytorch.org/blog/accelerating-2k-scale-pre-training-up-to-1-28x-with-torchao-mxfp8-and-torchtitan-on-crusoe-b200-cluster/). - -![MXFP8 vs BF16 Training Loss Curves](../../../assets/images/mxfp8_with_loss.png) - -*Training loss curves over 3,000 steps showing MXFP8 achieves equivalent convergence to bfloat16 baseline.* +Training runs on 64 node GB200 cluster with TorchTitan Llama4 Scout show that MXFP8 MoE training has equivalent convergence to bfloat16 training baseline over 3,000 steps. In fact, it finishes with slightly *lower* loss than bfloat16! This is consistent with our scaling experiments with [MXFP8 training for dense models](https://pytorch.org/blog/accelerating-2k-scale-pre-training-up-to-1-28x-with-torchao-mxfp8-and-torchtitan-on-crusoe-b200-cluster/). Training and model configurations for this run: - Model: Llama4 Scout @@ -199,6 +305,34 @@ Training and model configurations for this run: - `mxfp8` applied to routed experts computation (grouped GEMMs) - `mxfp8` applied to all linear layers except: `output`, `router.gate`, `attention.wk`, `attention.wv` (Wk and Wv too small to benefit from mxfp8) +#### Dense model convergence + +A deterministic 3,000-step comparison on C4 shows that Llama 3 8B with +32x32 MXFP8 weights closely tracks the BF16 baseline over 196.6M tokens. The +lower panel shows the difference between the 50-step mean losses, making the +small numerical divergence visible rather than implying bitwise-identical +training. + +![Llama 3 8B BF16 and MXFP8 32x32 training loss on C4](../../../../assets/images/mxfp8_32x32_vs_bf16_loss.png) + +*Training loss over 3,000 steps; faint lines are per-step values and bold lines +are 50-step moving averages.* + +Training and model configurations for this run: + +- Model: Llama 3 8B +- Dataset: C4 +- Hardware: 4x NVIDIA GB300 +- Training: 3,000 steps, 65,536 tokens/step, 196.6M tokens total +- Sequence length: 8192 +- Learning rate: 3e-4 +- LR scheduler warmup steps: 600 +- Parallelism: FSDP=4 +- Activation checkpointing: selective +- Seed: 42 with deterministic mode enabled +- `torch.compile` enabled +- MXFP8 weights use 32x32 scaling; BF16 is the baseline + ### Composability For distributed training, MXFP8 is compatible with: - `torch.compile` @@ -210,7 +344,6 @@ All distributed communication for MXFP8 training is currently done in high preci ### Known Limitations - Currently in prototype stage - no BC guarantees. - Requires torch nightly - important bug fixes have landed since 2.9.1 -- For GB200s, requires building torchao from source ### Additional Resources diff --git a/torchtitan/experiments/forge/__init__.py b/torchtitan/components/quantization/mxfp8/__init__.py similarity index 53% rename from torchtitan/experiments/forge/__init__.py rename to torchtitan/components/quantization/mxfp8/__init__.py index e4eecbcd35..47bb4472a1 100644 --- a/torchtitan/experiments/forge/__init__.py +++ b/torchtitan/components/quantization/mxfp8/__init__.py @@ -4,6 +4,11 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -from .engine import ForgeEngine +from .converter import MXFP8GroupedExpertsConverter, MXFP8Linear, MXFP8LinearConverter -__all__ = ["ForgeEngine"] + +__all__ = [ + "MXFP8GroupedExpertsConverter", + "MXFP8Linear", + "MXFP8LinearConverter", +] diff --git a/torchtitan/components/quantization/mx.py b/torchtitan/components/quantization/mxfp8/converter.py similarity index 55% rename from torchtitan/components/quantization/mx.py rename to torchtitan/components/quantization/mxfp8/converter.py index e6100d9c05..37fa7b8caf 100644 --- a/torchtitan/components/quantization/mx.py +++ b/torchtitan/components/quantization/mxfp8/converter.py @@ -11,40 +11,25 @@ from torchtitan.components.quantization import QuantizationConverter from torchtitan.models.common.linear import Linear from torchtitan.models.common.moe import GroupedExperts -from torchtitan.protocols.module import Module from torchtitan.tools.logging import logger from torchtitan.tools.utils import has_cuda_capability -from .utils import swap_token_dispatcher +from ..utils import swap_token_dispatcher -try: - from torchao.prototype.moe_training.mxfp8_linear import ( - MXFP8Linear as TorchAOMXFP8Linear, - ) - - class MXFP8Linear(TorchAOMXFP8Linear, Module): - """Inherits from Module (not Linear) to satisfy the Module protocol - (init_states, _param_init) while avoiding MRO conflicts with - Linear.__init__. Config still inherits from Linear.Config for - field compatibility. - """ - - @dataclass(kw_only=True, slots=True) - class Config(Linear.Config): - """Drop-in replacement for Linear.Config that builds MXFP8Linear.""" - - pass - - def __init__(self, config: Config): - TorchAOMXFP8Linear.__init__( - self, - config.in_features, - config.out_features, - bias=config.bias, - ) +_mxfp8_linear_import_error: ImportError | None = None -except ImportError: +try: + # Nothing about the class itself can fail here. What raises is two levels + # down: linear.py and tensor.py import torchao's mxfp8 cast kernels at + # module scope, and triton_to_mxfp8_32x32_swizzle_dim0_qdata_dim01_scale + # is newer than any torchao release. Catching it keeps + # ``import torchtitan.components.quantization`` working for float8 and + # nvfp4 users, and defers the error to whoever builds this converter. + from .linear import MXFP8Linear + +except ImportError as import_error: MXFP8Linear = None + _mxfp8_linear_import_error = import_error class MXFP8LinearConverter(QuantizationConverter): @@ -58,43 +43,107 @@ class Config(QuantizationConverter.Config): Only Linear.Config entries whose FQN contains a match are converted. If empty, all Linear modules are converted. """ + linears_saving_inputs_for_backward_in_mxfp8: list[str] = field( + default_factory=list + ) + """FQN substrings selecting linears that save inputs in MXFP8 for backward. + + A linear can save either its BF16 input or a columnwise MXFP8 input for + the backward pass. + + Without activation checkpointing, if the preceding operation already + saves its BF16 output for backward, as flash attention does, that tensor + is also available as this linear's input. Saving another MXFP8 + operands would increase memory usage, so this linear should save + BF16. If no other operation retains the BF16 input, saving MXFP8 reduces + activation memory and avoids columnwise quantization during backward. + + With full activation checkpointing, saved tensors from the original + forward are discarded and reconstructed during backward. Today, a + linear selected here produces its columnwise MXFP8 input in both the + original forward and recomputation, even though the original result is + discarded. An ideal checkpoint-aware policy could produce it only + during recomputation, but distinguishing those executions would add + complexity to the linear and its autograd contract. We intentionally + apply the same policy to both. + + More granular ``torch.remat`` policies add further save-versus-recompute + choices, so the optimal format depends on both model activation + ownership and the activation-checkpointing policy. BF16 is therefore + the conservative default, and users can opt selected modules into MXFP8 + with this list. + """ + + def __post_init__(self) -> None: + if any(not fqn for fqn in self.linears_saving_inputs_for_backward_in_mxfp8): + raise ValueError( + "MXFP8 linears_saving_inputs_for_backward_in_mxfp8 cannot " + "contain an empty FQN selector." + ) def __init__(self, config: Config): self.config = config if MXFP8Linear is None: raise ImportError( - "torchao is not installed. Please install it to use MXFP8 linear layers." - ) + "MXFP8 linear layers need torchao's 32x32 swizzled cast " + "kernels, added in pytorch/ao#4777 and not in any release up " + "to v0.18.0. Install a torchao that contains it." + ) from _mxfp8_linear_import_error if not has_cuda_capability(10, 0): raise ValueError("MXFP8 is only supported on SM100 or later architectures") - if not self.config.model_compile_enabled: - logger.warning( - "torch.compile enablement is required for highest performance " - "of MXFP8 dynamic quantization." - ) - def convert(self, model_config): assert MXFP8Linear is not None fqns = self.config.fqns - for fqn, config, parent, attr in model_config.traverse(Linear.Config): - if not fqns or any(target_fqn in fqn for target_fqn in fqns): - new_config = MXFP8Linear.Config( - in_features=config.in_features, - out_features=config.out_features, - bias=config.bias, - param_init=config.param_init, - ) - if parent is None: - model_config = new_config - elif isinstance(parent, list): - parent[attr] = new_config - else: - setattr(parent, attr, new_config) - - logger.info("Converted Linear layers to MXFP8Linear") + targets = [ + entry + for entry in model_config.traverse(Linear.Config) + if not fqns or any(target_fqn in entry[0] for target_fqn in fqns) + ] + + selectors = self.config.linears_saving_inputs_for_backward_in_mxfp8 + target_fqns = [fqn for fqn, _config, _parent, _attr in targets] + unmatched_fqn_selectors = { + selector + for selector in selectors + if not any(selector in fqn for fqn in target_fqns) + } + if unmatched_fqn_selectors: + raise ValueError( + "MXFP8 linears_saving_inputs_for_backward_in_mxfp8 selectors " + "did not match any converted Linear.Config: " + f"{sorted(unmatched_fqn_selectors)}." + ) + + mxfp8_fqns = { + fqn for fqn in target_fqns if any(selector in fqn for selector in selectors) + } + for fqn, config, parent, attr in targets: + new_config = MXFP8Linear.Config( + in_features=config.in_features, + out_features=config.out_features, + bias=config.bias, + param_init=config.param_init, + input_activation_format_for_backward=( + "mxfp8" if fqn in mxfp8_fqns else "bf16" + ), + ) + if parent is None: + model_config = new_config + elif isinstance(parent, list): + parent[attr] = new_config + else: + setattr(parent, attr, new_config) + + num_mxfp8 = len(mxfp8_fqns) + num_bf16 = len(targets) - num_mxfp8 + logger.info( + "Converted Linear layers to MXFP8Linear with saved input activation " + f"formats: {num_bf16} bf16, {num_mxfp8} mxfp8" + ) + logger.debug(f"Linears saving MXFP8 input activations: {sorted(mxfp8_fqns)}") return model_config diff --git a/torchtitan/components/quantization/mxfp8/linear.py b/torchtitan/components/quantization/mxfp8/linear.py new file mode 100644 index 0000000000..cf875de346 --- /dev/null +++ b/torchtitan/components/quantization/mxfp8/linear.py @@ -0,0 +1,411 @@ +# 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. + +"""MXFP8 linear training with FSDP-managed 32x32 weight caches. + +Tensor shape suffixes: + M: flattened token rows + N: output features + K: input features +""" + +from dataclasses import dataclass +from typing import Literal + +import spmd_types as spmd +import torch +import torch.nn.functional as F +from torch import nn +from torch.autograd.function import once_differentiable + +from torchao.prototype.mx_formats.kernels import ( + mxfp8_quantize_cuda, + triton_mx_block_rearrange, +) + +from torchtitan.distributed.utils import get_spmd_backend +from torchtitan.models.common.linear import Linear + +from .._fsdp_tensor import _UnshardedFSDPTensor +from .tensor import ( + _LinearShardedTensorWithMXFP8Compute, + _MXFP8_BLOCK_SIZE, + _quantize_mxfp8_weight, +) + + +__all__ = ["InputActivationFormatForBackward", "MXFP8Linear"] + +# Activation and gradient quantization takes a scaling mode; the 32x32 weight +# cast hardcodes RCEIL. Pin the two to match, so both operands of a GEMM round +# their E8M0 scales the same way. This is TorchAO's current default too, but +# relying on that would let a default change silently desync them. +_MXFP8_SCALING_MODE = "rceil" + +InputActivationFormatForBackward = Literal["bf16", "mxfp8"] +_INPUT_ACTIVATION_FORMATS_FOR_BACKWARD = ("bf16", "mxfp8") + + +def _pad_rows(x_MK: torch.Tensor) -> tuple[torch.Tensor, int]: + num_rows = x_MK.shape[0] + num_padded_rows = ( + (num_rows + _MXFP8_BLOCK_SIZE - 1) // _MXFP8_BLOCK_SIZE + ) * _MXFP8_BLOCK_SIZE + if num_padded_rows == num_rows: + return x_MK, num_rows + return F.pad(x_MK, (0, 0, 0, num_padded_rows - num_rows)), num_rows + + +# Adapted from torchao.prototype.moe_training.mxfp8_linear.mx_mm. This variant +# lives in TorchTitan so its autograd state and weight cache can integrate with +# FSDP and other parallelisms. +@torch._dynamo.allow_in_graph +class _MXFP8LinearFunction(torch.autograd.Function): + @staticmethod + # pyrefly: ignore [bad-override] + def forward( + ctx, + x: torch.Tensor, + weight_NK: torch.Tensor, + weight_qdata_fprop_KN: torch.Tensor, + weight_scale_fprop_swizzled: torch.Tensor, + weight_qdata_dgrad_NK: torch.Tensor, + weight_scale_dgrad_swizzled: torch.Tensor, + bias_N: torch.Tensor | None, + input_activation_format_for_backward: InputActivationFormatForBackward, + ) -> torch.Tensor: + # FPROP always consumes rowwise MXFP8. WGRAD can either retain the + # original BF16 input and quantize it columnwise in backward, or retain + # a columnwise MXFP8 operands produced in forward. The former is + # memory-safe when another operation already keeps BF16 x alive; the + # latter reduces storage for a single-consumer input at the cost of an + # extra cached operands when BF16 x is retained elsewhere. Under + # full activation checkpointing, the selected state is created by + # recompute. + if x.dtype != torch.bfloat16 or weight_NK.dtype != torch.bfloat16: + raise ValueError( + "MXFP8Linear requires BF16 activations and weights; " + f"got activation dtype {x.dtype} and weight dtype {weight_NK.dtype}." + ) + if bias_N is not None and bias_N.dtype != torch.bfloat16: + raise ValueError( + f"MXFP8Linear requires a BF16 bias; got bias dtype {bias_N.dtype}." + ) + if x.shape[-1] != weight_NK.shape[1]: + raise ValueError( + "MXFP8Linear activation and weight contraction dimensions must " + f"match; got {x.shape[-1]} and {weight_NK.shape[1]}." + ) + for name, value in ( + ("local in_features", weight_NK.shape[1]), + ("local out_features", weight_NK.shape[0]), + ): + if value % _MXFP8_BLOCK_SIZE: + raise ValueError( + f"MXFP8Linear requires {name} divisible by " + f"{_MXFP8_BLOCK_SIZE}; got {value}." + ) + + input_shape = x.shape + x_MK, num_rows = _pad_rows(x.reshape(-1, input_shape[-1]).contiguous()) + requires_wgrad = ctx.needs_input_grad[1] + quantize_wgrad_input_in_forward = ( + requires_wgrad and input_activation_format_for_backward == "mxfp8" + ) + + # The save format controls both computation and saved state. BF16 mode + # requests only the rowwise FPROP operand here; backward produces the + # columnwise WGRAD operand from the saved BF16 input. + # TODO(anijain2305): torchao's mxfp8_quantize_2d_{1x32,32x1}_cutedsl + # fuse the cast and the scale swizzle into one kernel, replacing this + # call plus the triton_mx_block_rearrange below. Measured 2.4-3.2x + # faster than the pair on a GB200, bitwise identical on both outputs. + # Three things to settle before switching: + # - They require the token count to be a multiple of 128, where this + # path needs only 32. The 32 is the MX scaling granularity, and the + # 128 is the tcgen05 scale-tile height that scaled_mm wants either + # way -- splitting the two kernels is what lets + # triton_mx_block_rearrange pad the *scales* up to 128 rows and + # leave the token count alone. Fusing pushes that padding onto the + # activations, so a 64-token microbatch would run the quantizer and + # the GEMM over 128 rows. The speedup above was measured on shapes + # that already divide 128 and should not be assumed to hold once + # small token counts pay for the extra rows. + # - They need nvidia-cutlass-dsl and apache-tvm-ffi, which torchao + # does not depend on: MXFP8 dense linears work without them today, + # and switching would make them mandatory for every MXFP8 user. + # - The usable cutlass-dsl range is narrow. torchao's README asks for + # 4.5.2; 4.6.0 changed the nvvm.cvt_packfloat* builders and breaks + # torchao's CuTeDSL kernels outright. + # Their availability check also raises from inside the kernel, so a + # missing package would surface on the first forward. Gate it in + # MXFP8LinearConverter.__init__ instead, beside the torchao check. + x_qdata_row_MK, x_qdata_col_MK, x_scale_row, x_scale_col = mxfp8_quantize_cuda( + x_MK, + rowwise=True, + colwise=quantize_wgrad_input_in_forward, + scaling_mode=_MXFP8_SCALING_MODE, + ) + x_scale_row = triton_mx_block_rearrange(x_scale_row) + if quantize_wgrad_input_in_forward: + x_scale_col = triton_mx_block_rearrange(x_scale_col) + + # The 32x32 weight quantizer returns both qdata/scale pairs ready for + # this exact BlockWise1x32 and SWIZZLE_32_4_4 B-operand contract. + output_MN = F.scaled_mm( + x_qdata_row_MK, + weight_qdata_fprop_KN, + scale_a=x_scale_row, + scale_recipe_a=F.ScalingType.BlockWise1x32, + scale_b=weight_scale_fprop_swizzled, + scale_recipe_b=F.ScalingType.BlockWise1x32, + swizzle_a=F.SwizzleType.SWIZZLE_32_4_4, + swizzle_b=F.SwizzleType.SWIZZLE_32_4_4, + bias=bias_N, + output_dtype=torch.bfloat16, + ) + + # Save exactly one input-activation operands for WGRAD. BF16 mode + # keeps the original tensor and builds the columnwise operand in + # backward. MXFP8 mode keeps the columnwise qdata and scales produced + # above. FPROP and DGRAD share the same weight qdata allocation. + # An unsharded tensor's storage is FSDP's to free at reshard and refill + # before backward, so save the wrapper and read the operands off it + # then. Anything else carries no operands to refill, so save them. + has_unsharded_tensor = isinstance(weight_NK, _UnshardedFSDPTensor) + saved_weight_tensors = ( + (weight_NK,) + if has_unsharded_tensor + else (weight_qdata_dgrad_NK, weight_scale_dgrad_swizzled) + ) + if requires_wgrad and input_activation_format_for_backward == "bf16": + ctx.save_for_backward(x, *saved_weight_tensors) + else: + ctx.save_for_backward(x_qdata_col_MK, x_scale_col, *saved_weight_tensors) + ctx.has_unsharded_tensor = has_unsharded_tensor + ctx.input_shape = input_shape + ctx.num_rows = num_rows + ctx.requires_dgrad = ctx.needs_input_grad[0] + ctx.requires_wgrad = requires_wgrad + ctx.input_activation_format_for_backward = input_activation_format_for_backward + ctx.has_bias = bias_N is not None + + return output_MN[:num_rows].reshape(*input_shape[:-1], weight_NK.shape[0]) + + @staticmethod + @once_differentiable + # pyrefly: ignore [bad-override] + def backward(ctx, grad_output: torch.Tensor): + # WGRAD consumes either the saved columnwise activation pair or a pair + # rebuilt from the saved BF16 input. DGRAD consumes the weight pair. + x_hp = None + x_qdata_col_MK = None + x_scale_col = None + saved_tensors = ctx.saved_tensors + if ctx.requires_wgrad and ctx.input_activation_format_for_backward == "bf16": + x_hp = saved_tensors[0] + saved_weight_tensors = saved_tensors[1:] + else: + x_qdata_col_MK, x_scale_col = saved_tensors[:2] + saved_weight_tensors = saved_tensors[2:] + + if ctx.has_unsharded_tensor: + (weight_NK,) = saved_weight_tensors + if not isinstance(weight_NK, _UnshardedFSDPTensor): + raise RuntimeError("FSDP restored an incompatible MXFP8 weight") + operands = weight_NK.operands + weight_qdata_dgrad_NK = operands.weight_qdata_dgrad_NK + weight_scale_dgrad_swizzled = operands.weight_scale_dgrad_swizzled + else: + weight_qdata_dgrad_NK, weight_scale_dgrad_swizzled = saved_weight_tensors + + grad_output_MN = grad_output.contiguous().reshape(-1, grad_output.shape[-1]) + grad_bias_N = grad_output_MN.sum(dim=0) if ctx.has_bias else None + + grad_input = None + grad_weight_NK = None + if ctx.requires_dgrad or ctx.requires_wgrad: + padded_grad_output_MN, _ = _pad_rows(grad_output_MN) + ( + grad_output_row_MN, + grad_output_col_MN, + grad_output_row_scales, + grad_output_col_scales, + ) = mxfp8_quantize_cuda( + padded_grad_output_MN, + rowwise=ctx.requires_dgrad, + colwise=ctx.requires_wgrad, + scaling_mode=_MXFP8_SCALING_MODE, + ) + + if ctx.requires_dgrad: + grad_output_row_scales = triton_mx_block_rearrange( + grad_output_row_scales + ) + grad_input_MK = F.scaled_mm( + grad_output_row_MN, + weight_qdata_dgrad_NK, + scale_a=grad_output_row_scales, + scale_recipe_a=F.ScalingType.BlockWise1x32, + scale_b=weight_scale_dgrad_swizzled, + scale_recipe_b=F.ScalingType.BlockWise1x32, + swizzle_a=F.SwizzleType.SWIZZLE_32_4_4, + swizzle_b=F.SwizzleType.SWIZZLE_32_4_4, + output_dtype=torch.bfloat16, + ) + grad_input = grad_input_MK[: ctx.num_rows].reshape(ctx.input_shape) + + if ctx.requires_wgrad: + if ctx.input_activation_format_for_backward == "bf16": + assert x_hp is not None + x_MK, _ = _pad_rows( + x_hp.reshape(-1, ctx.input_shape[-1]).contiguous() + ) + _, x_qdata_col_MK, _, x_scale_col = mxfp8_quantize_cuda( + x_MK, + rowwise=False, + colwise=True, + scaling_mode=_MXFP8_SCALING_MODE, + ) + x_scale_col = triton_mx_block_rearrange(x_scale_col) + + assert x_qdata_col_MK is not None + assert x_scale_col is not None + grad_output_col_scales = triton_mx_block_rearrange( + grad_output_col_scales + ) + grad_weight_NK = F.scaled_mm( + grad_output_col_MN.t(), + x_qdata_col_MK, + scale_a=grad_output_col_scales, + scale_recipe_a=F.ScalingType.BlockWise1x32, + scale_b=x_scale_col, + scale_recipe_b=F.ScalingType.BlockWise1x32, + swizzle_a=F.SwizzleType.SWIZZLE_32_4_4, + swizzle_b=F.SwizzleType.SWIZZLE_32_4_4, + output_dtype=torch.bfloat16, + ) + + return grad_input, grad_weight_NK, None, None, None, None, grad_bias_N, None + + +# Marks the function local-only so SPMD type checking can propagate through +# an autograd function it cannot see into. +# TODO(anijain2305, pianpwk): drop this once register_local_autograd_function +# is removed tree-wide. nvfp4 and qwen3_5's gdn still rely on the same +# registration, so it has to go everywhere at once. +spmd.register_local_autograd_function(_MXFP8LinearFunction) + + +class MXFP8Linear(Linear): + """Linear using 1D activations and cached 32x32 weight quantization.""" + + @dataclass(kw_only=True, slots=True) + class Config(Linear.Config): + """Drop-in replacement for ``Linear.Config``.""" + + input_activation_format_for_backward: InputActivationFormatForBackward = "bf16" + """Format used to save the input activation needed by WGRAD. + + ``"bf16"`` saves the original input and quantizes it columnwise during + backward. ``"mxfp8"`` produces the columnwise operands during + forward and saves its qdata and scales for backward. + """ + + def __post_init__(self) -> None: + if ( + self.input_activation_format_for_backward + not in _INPUT_ACTIVATION_FORMATS_FOR_BACKWARD + ): + raise ValueError( + "MXFP8 input_activation_format_for_backward must be one of " + f"{_INPUT_ACTIVATION_FORMATS_FOR_BACKWARD}; got " + f"{self.input_activation_format_for_backward!r}." + ) + for name in ("in_features", "out_features"): + value = getattr(self, name) + if value % _MXFP8_BLOCK_SIZE: + raise ValueError( + f"MXFP8 requires {name} divisible by {_MXFP8_BLOCK_SIZE}; " + f"got {name}={value}." + ) + + def build(self, **kwargs): + # The MXFP8 matmul is an opaque autograd function, so DTensor has + # no sharding strategy for it. Under partial_dtensor with TP, + # propagation reaches into the storage-free unsharded tensor and + # fails; making it work needs local_map plus hand-declared input + # and input-gradient placements. spmd_types instead annotates the + # function itself (see register_local_autograd_function above), so + # the stock Linear sharding config suffices there. Reject the + # backend rather than carry a second sharding path for it. + if get_spmd_backend() == "partial_dtensor": + raise ValueError( + "MXFP8Linear requires parallelism.spmd_backend=" + "'spmd_types'; got 'partial_dtensor'. The MXFP8 matmul is " + "an opaque autograd function with no DTensor sharding " + "rule, so tensor parallelism cannot propagate through it." + ) + return Linear.Config.build(self, **kwargs) + + def __init__(self, config: Config): + super().__init__(config) + self.input_activation_format_for_backward = ( + config.input_activation_format_for_backward + ) + # Install the unsharded-tensor wrapper up front so no caller has to + # remember to do it. The wrapper is inert until a data parallel + # implementation drives its unshard lifecycle: until then it just holds + # the BF16 weight, and forward rejects it. + self.weight = nn.Parameter( + _LinearShardedTensorWithMXFP8Compute(self.weight.data), + requires_grad=self.weight.requires_grad, + ) + + def forward(self, input: torch.Tensor) -> torch.Tensor: + # Always a plain tensor: the weight is only re-wrapped as a DTensor on + # a non-data-parallel mesh under partial_dtensor, which Config.build + # rejects, and spmd_types carries TP and EP as annotations instead. + weight_NK = self.weight + # __init__ installs a _LinearShardedTensorWithMXFP8Compute, but that is + # not what forward usually sees. Under FSDP the post-all-gather hook has + # already replaced it for this unshard lifetime with the storage-free + # _UnshardedFSDPTensor holding the quantized operands, so the weight + # arrives here already quantized and the type identifies which state we + # are in. + if isinstance(weight_NK, _UnshardedFSDPTensor): + operands = weight_NK.operands + else: + # No data parallel implementation owns this weight's lifecycle, so + # it still holds high-precision storage and the operands are built + # per invocation. Eager FSDP2 always installs an unsharded tensor, but + # GraphTrainer under the spmd_types backend does not: its runtime + # hands forward a plain annotated local tensor, so the wrapper + # SimpleFSDP's parametrization built never reaches here. Quantize + # the storage rather than the wrapper, which the kernels cannot + # consume; ``weight_NK`` itself stays wrapped so autograd returns + # the gradient to the parameter. + with torch.no_grad(): + operands = _quantize_mxfp8_weight( + weight_NK._tensor + if isinstance(weight_NK, _LinearShardedTensorWithMXFP8Compute) + else weight_NK + ) + # Nothing caches this across calls, so a frozen weight is + # requantized on every forward. Training pays that anyway, since + # the weight changes each optimizer step; inference does not. + # TODO(anijain2305): key the operands on the parameter's + # version counter so a frozen weight is quantized once. + return _MXFP8LinearFunction.apply( + input, + weight_NK, + operands.weight_qdata_fprop_KN, + operands.weight_scale_fprop_swizzled, + operands.weight_qdata_dgrad_NK, + operands.weight_scale_dgrad_swizzled, + self.bias, + self.input_activation_format_for_backward, + ) diff --git a/torchtitan/components/quantization/mxfp8/tensor.py b/torchtitan/components/quantization/mxfp8/tensor.py new file mode 100644 index 0000000000..57ba72151b --- /dev/null +++ b/torchtitan/components/quantization/mxfp8/tensor.py @@ -0,0 +1,103 @@ +# 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. + +"""MXFP8 specialization of the generic FSDP unsharded-tensor lifecycle.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from torchao.prototype.mx_formats.kernels import ( + triton_to_mxfp8_32x32_swizzle_dim0_qdata_dim01_scale, +) + +from .._fsdp_tensor import _ShardedFSDPTensor + + +# Everything here is internal to the MXFP8 component; nothing is re-exported. +__all__: list[str] = [] + +# One E8M0 scale per 32 elements along the scaled axis, per the OCP +# microscaling spec. Weight quantization uses a *square* 32x32 tile -- +# side equal to the block size is the only shape that is a valid MX +# group along both axes at once, which is what makes it transpose- +# invariant and lets FPROP and DGRAD share one qdata allocation. +_MXFP8_BLOCK_SIZE = 32 + + +@dataclass(frozen=True, slots=True) +class _MXFP8LinearOperands: + """The independent MXFP8 tensors owned by one FSDP unshard lifetime. + + Everything here is quantized the same way: square 32x32 tiles, E4M3 qdata + and one E8M0 scale per tile. ``swizzled`` on the scales names their memory + layout, not their format -- they are pre-arranged into the blocked grid + ``scaled_mm`` wants for ``SwizzleType.SWIZZLE_32_4_4``, so no rearrange is + needed at the GEMM. Square tiles are transpose-invariant, so FPROP and + DGRAD share one qdata allocation and differ only in scale layout. + """ + + weight_qdata_dgrad_NK: torch.Tensor # noqa: N815 + weight_scale_fprop_swizzled: torch.Tensor + weight_scale_dgrad_swizzled: torch.Tensor + + @property + def weight_qdata_fprop_KN(self) -> torch.Tensor: # noqa: N802 + return self.weight_qdata_dgrad_NK.t() + + +def _quantize_mxfp8_weight(weight_NK: torch.Tensor) -> _MXFP8LinearOperands: + """Quantize a BF16 weight using fixed square 32x32 scale tiles.""" + if weight_NK.ndim != 2: + raise ValueError( + "MXFP8 32x32 weight quantization requires a 2D weight, " + f"got {weight_NK.ndim} dimensions." + ) + if weight_NK.dtype != torch.bfloat16: + raise ValueError( + "MXFP8 32x32 weight quantization requires BF16 weights, " + f"got {weight_NK.dtype}." + ) + if any(size % _MXFP8_BLOCK_SIZE for size in weight_NK.shape): + raise ValueError( + "MXFP8 32x32 weight quantization requires both matrix dimensions " + f"divisible by {_MXFP8_BLOCK_SIZE}, got {tuple(weight_NK.shape)}." + ) + ( + weight_qdata_dgrad_NK, + weight_scale_fprop_swizzled, + weight_scale_dgrad_swizzled, + ) = triton_to_mxfp8_32x32_swizzle_dim0_qdata_dim01_scale(weight_NK) + return _MXFP8LinearOperands( + weight_qdata_dgrad_NK=weight_qdata_dgrad_NK, + weight_scale_fprop_swizzled=weight_scale_fprop_swizzled, + weight_scale_dgrad_swizzled=weight_scale_dgrad_swizzled, + ) + + +class _LinearShardedTensorWithMXFP8Compute(_ShardedFSDPTensor): + """The persistent BF16 linear parameter; quantizes to MXFP8 on unshard. + + This is the sharded state only. FSDP shards, all-gathers, reduces + gradients into, and checkpoints this BF16 parameter. The MXFP8 operands it + produces live on a ``_UnshardedFSDPTensor`` for one unshard lifetime; that + holder is generic, so quantization is the only thing a format supplies. + """ + + def _build_operands( + self, + logical_tensor: torch.Tensor, + out: _MXFP8LinearOperands | None = None, + ) -> _MXFP8LinearOperands: + operands = _quantize_mxfp8_weight(logical_tensor) + if out is None: + return operands + out.weight_qdata_dgrad_NK.copy_(operands.weight_qdata_dgrad_NK) + out.weight_scale_fprop_swizzled.copy_(operands.weight_scale_fprop_swizzled) + out.weight_scale_dgrad_swizzled.copy_(operands.weight_scale_dgrad_swizzled) + return out diff --git a/torchtitan/components/quantization/utils.py b/torchtitan/components/quantization/utils.py index 0de363dfa8..837b6889bd 100644 --- a/torchtitan/components/quantization/utils.py +++ b/torchtitan/components/quantization/utils.py @@ -70,7 +70,10 @@ def has_quantization(model_config) -> bool: _float8_experts_cache, Float8Linear, ) - from torchtitan.components.quantization.mx import _mxfp8_experts_cache, MXFP8Linear + from torchtitan.components.quantization.mxfp8.converter import ( + _mxfp8_experts_cache, + MXFP8Linear, + ) from torchtitan.components.quantization.nvfp4 import NVFP4Linear quant_linear_types: list[type] = [] diff --git a/torchtitan/distributed/cudagraph.py b/torchtitan/distributed/cudagraph.py index b9dafeb7d6..fe355b6d50 100644 --- a/torchtitan/distributed/cudagraph.py +++ b/torchtitan/distributed/cudagraph.py @@ -6,15 +6,12 @@ """Lightweight CUDA graph wrapper for training steps.""" -import gzip -import json import warnings from collections.abc import Callable, Sequence from dataclasses import dataclass from typing import Any, cast import torch -from torch.cuda._annotate_cuda_graph_trace import annotate_trace from torch.cuda._graph_annotations import get_kernel_annotations from torch.nn.attention.flex_attention import BlockMask from torch.utils import _pytree as pytree @@ -193,23 +190,6 @@ def get_cudagraph_annotations() -> dict[int, list[Any]]: return _manager.all_annotations -def cudagraph_annotate_trace_post_processor(trace_path: str) -> None: - """Post-process a profiler trace with captured CUDA graph annotations.""" - annotations = get_cudagraph_annotations() - if not annotations: - return - - open_trace = gzip.open if trace_path.endswith(".gz") else open - with open_trace(trace_path, "rt") as trace_file: - trace = json.load(trace_file) - - count = annotate_trace(trace, annotations) - if count > 0: - with open_trace(trace_path, "wt") as trace_file: - json.dump(trace, trace_file) - logger.info(f"Annotated {count} CUDA graph kernel events in profiler trace") - - class CUDAGraphWrapper: """Wrap a callable with CUDA graph capture and replay. diff --git a/torchtitan/experiments/README.md b/torchtitan/experiments/README.md index 30c4142697..a61b54168a 100644 --- a/torchtitan/experiments/README.md +++ b/torchtitan/experiments/README.md @@ -28,4 +28,3 @@ We provide this `experiments/` folder to host experiments that add significant v | [torchft](./torchft/) | [![TorchFT 8 GPU Integration Tests](https://github.com/pytorch/torchtitan/actions/workflows/integration_test_8gpu_torchft.yaml/badge.svg?branch=main)](https://github.com/pytorch/torchtitan/actions/workflows/integration_test_8gpu_torchft.yaml?query=branch%3Amain) | [@tushar00jain](https://github.com/tushar00jain) [@fegin](https://github.com/fegin) | | [transformers_modeling_backend](./transformers_modeling_backend/) | [![Transformers modeling backend 8 GPU Integration Tests](https://github.com/pytorch/torchtitan/actions/workflows/integration_test_8gpu_transformers_modeling_backend.yaml/badge.svg?branch=main)](https://github.com/pytorch/torchtitan/actions/workflows/integration_test_8gpu_transformers_modeling_backend.yaml?query=branch%3Amain) | [@3outeille](https://github.com/3outeille) [@mreso](https://github.com/mreso) | | [rl](./rl/) | [![RL 8 GPU Integration Tests](https://github.com/pytorch/torchtitan/actions/workflows/integration_test_8gpu_rl.yaml/badge.svg?branch=main)](https://github.com/pytorch/torchtitan/actions/workflows/integration_test_8gpu_rl.yaml?query=branch%3Amain) | [@wwwjn](https://github.com/wwwjn) [@felipemello1](https://github.com/felipemello1) | -| [forge](./forge/) | TBA | [@felipemello1](https://github.com/felipemello1) | diff --git a/torchtitan/experiments/forge/README.md b/torchtitan/experiments/forge/README.md deleted file mode 100644 index a99b8b476b..0000000000 --- a/torchtitan/experiments/forge/README.md +++ /dev/null @@ -1,12 +0,0 @@ -## `ForgeEngine` - -The `forge` folder contains a lightweight training engine that serves as a streamlined subset of the `Trainer` class from [torchtitan/train.py](/torchtitan/train.py). This engine provides only the essential constructor method, making it highly flexible for various downstream applications. - -The [`ForgeEngine`](engine.py) takes a `ForgeEngine.Config` to -- Initialize an SPMD distributed training environment -- Construct and scale models via n-D parallelisms and meta-device initialization -- Provide necessary training components and utilities - -**Primary Use Case**: The engine is designed for building trainers in post-training workflows where multiple specialized components (trainer, generator, replay buffer, parameter server, etc.) work together. - -The [example_train.py](./example_train.py) demonstrates how to use `ForgeEngine` for pretraining, achieving the same functionality as [torchtitan/train.py](/torchtitan/train.py) (except for quantization or fault tolerance). diff --git a/torchtitan/experiments/forge/engine.py b/torchtitan/experiments/forge/engine.py deleted file mode 100644 index c830236fa7..0000000000 --- a/torchtitan/experiments/forge/engine.py +++ /dev/null @@ -1,306 +0,0 @@ -# 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 os -from collections.abc import Generator -from dataclasses import asdict, dataclass, field -from typing import Any - -import torch -from torch.distributed.elastic.multiprocessing.errors import record - -from torchtitan.components.checkpointer import CheckpointManager -from torchtitan.components.loss import LossFunction -from torchtitan.components.optimizer import LRSchedulersContainer, OptimizersContainer -from torchtitan.config import Configurable, TORCH_DTYPE_MAP -from torchtitan.config.configs import ( - CommConfig, - CompileConfig, - DebugConfig, - ParallelismConfig, - TrainingConfig, -) -from torchtitan.distributed import ParallelDims, utils as dist_utils -from torchtitan.distributed.activation_checkpoint import ( - ActivationCheckpointingConfig, - MemoryBudgetAC, - SelectiveAC, -) -from torchtitan.protocols import BaseModel -from torchtitan.protocols.model_spec import ModelSpec -from torchtitan.tools import utils - - -class ForgeEngine(torch.distributed.checkpoint.stateful.Stateful, Configurable): - @dataclass(kw_only=True, slots=True) - class Config(Configurable.Config): - hf_assets_path: str = "./tests/assets/tokenizer" - dump_folder: str = "./outputs" - model_spec: ModelSpec = field(default_factory=ModelSpec) - optimizer: OptimizersContainer.Config = field( - default_factory=OptimizersContainer.Config - ) - lr_scheduler: LRSchedulersContainer.Config = field( - default_factory=LRSchedulersContainer.Config - ) - training: TrainingConfig = field(default_factory=TrainingConfig) - parallelism: ParallelismConfig = field(default_factory=ParallelismConfig) - checkpoint: CheckpointManager.Config = field( - default_factory=CheckpointManager.Config - ) - activation_checkpoint: ActivationCheckpointingConfig = field( - default_factory=SelectiveAC.Config - ) - compile: CompileConfig = field(default_factory=CompileConfig) - comm: CommConfig = field(default_factory=CommConfig) - debug: DebugConfig = field(default_factory=DebugConfig) - - def __post_init__(self): - if isinstance(self.activation_checkpoint, MemoryBudgetAC.Config) and not ( - self.compile.enable and "model" in self.compile.components - ): - raise ValueError( - "Memory budget activation checkpointing requires the model to be " - "compiled: set --compile.enable and include 'model' in " - "--compile.components." - ) - - def to_dict(self) -> dict[str, Any]: - return asdict(self) - - # core configs - config: Config - parallel_dims: ParallelDims - train_spec: ModelSpec - - # swappable training components in ModelSpec - model_parts: list[torch.nn.Module] - loss_fn: LossFunction - optimizers: OptimizersContainer - lr_schedulers: LRSchedulersContainer - - # non-swappable training components - checkpointer: CheckpointManager - - # runtime utilities - device: torch.device - gc_handler: utils.GarbageCollection - gradient_accumulation_steps: int - num_pp_microbatches: int - train_context: Generator[None, None, None] - pp_has_first_stage: bool - pp_has_last_stage: bool - - # Fields in ForgeEngine which are not in original Trainer - # for dataloading - dp_degree: int - dp_rank: int - # for logging - model_config: BaseModel.Config - num_flops_per_token: float - model_param_count: int - num_tokens_per_train_step: int - - # Enable debug tracing on failure: https://pytorch.org/docs/stable/elastic/errors.html - @record - def __init__(self, config: Config): - torch._C._log_api_usage_once("torchtitan.train") - - self.config = config - - device_module, device_type = utils.device_module, utils.device_type - self.device = utils.get_local_device() - # Device has to be set before creating TorchFT manager. - device_module.set_device(self.device) - - # init distributed and build meshes - dist_utils.init_distributed( - config.comm, - enable_cpu_backend=config.training.enable_cpu_offload, - ) - world_size = int(os.environ["WORLD_SIZE"]) - self.parallel_dims = parallel_dims = ParallelDims.from_config( - config.parallelism, world_size - ) - - if parallel_dims.dp_enabled: - batch_mesh = parallel_dims.get_mesh("batch") - dp_degree, dp_rank = batch_mesh.size(), batch_mesh.get_local_rank() - else: - dp_degree, dp_rank = 1, 0 - self.dp_degree, self.dp_rank = dp_degree, dp_rank - - # take control of garbage collection to avoid stragglers - self.gc_handler = utils.GarbageCollection( - gc_freq=config.training.gc_freq, debug=config.training.gc_debug - ) - - # Set random seed, and maybe enable deterministic mode - # (mainly for debugging, expect perf loss). - dist_utils.set_determinism( - parallel_dims, - self.device, - config.debug, - distinct_seed_mesh_dims=["pp"], # same as `torchtitan/train.py` - ) - self.train_spec = config.model_spec - - # build model (using meta init) - self.model_config = model_config = self.train_spec.model - # set the model args from training configs - model_config.update_from_config( - config=config, - ) - - with ( - torch.device("meta"), - utils.set_default_dtype(TORCH_DTYPE_MAP[config.training.dtype]), - ): - model = model_config.build() - - # calculate model size and flops per token - ( - self.model_param_count, - self.num_flops_per_token, - ) = model_config.get_nparams_and_flops( - model, config.training.max_context_length - ) - - # move sharded model to CPU/GPU and initialize weights via DTensor - if config.training.enable_cpu_offload: - init_device = "cpu" - buffer_device = device_type - else: - init_device = device_type - buffer_device = None - - self.loss_fn = self.train_spec.loss.build( - config.compile, parallel_dims=parallel_dims - ) - - # Verify token budgets. - num_pp_microbatches = ( - config.parallelism.num_pp_microbatches if parallel_dims.pp_enabled else 1 - ) - self.num_pp_microbatches = num_pp_microbatches - self.num_tokens_per_train_step = config.training.num_tokens_per_train_step - if self.num_tokens_per_train_step < 0: - self.num_tokens_per_train_step = ( - config.training.num_tokens_per_microbatch_per_dp_rank - * self.num_pp_microbatches - * dp_degree - ) - if ( - self.num_tokens_per_train_step - % ( - config.training.num_tokens_per_microbatch_per_dp_rank - * self.num_pp_microbatches - * dp_degree - ) - != 0 - ): - raise ValueError( - "training.num_tokens_per_train_step " - f"({self.num_tokens_per_train_step}) must be divisible by the " - "number of tokens processed globally in one gradient accumulation " - "iteration " - f"({config.training.num_tokens_per_microbatch_per_dp_rank * self.num_pp_microbatches * dp_degree})." - ) - self.gradient_accumulation_steps = self.num_tokens_per_train_step // ( - config.training.num_tokens_per_microbatch_per_dp_rank - * self.num_pp_microbatches - * dp_degree - ) - - # apply parallelisms and initialization - if parallel_dims.pp_enabled: - if not self.train_spec.pipelining_fn: - raise RuntimeError( - f"Pipeline Parallel is enabled but {self.train_spec.name} " - f"does not support pipelining" - ) - - # apply both PT-D Pipeline Parallel and SPMD-style PT-D techniques - ( - self.pp_schedule, - self.model_parts, - self.pp_has_first_stage, - self.pp_has_last_stage, - ) = self.train_spec.pipelining_fn( - model, - parallel_dims=parallel_dims, - training=config.training, - parallelism=config.parallelism, - compile_config=config.compile, - ac_config=config.activation_checkpoint, - dump_folder=config.dump_folder, - device=self.device, - model_config=model_config, - parallelize_fn=self.train_spec.parallelize_fn, - loss_fn=self.loss_fn, - ) - # when PP is enabled, `model` obj is no longer used after this point, - # model_parts is used instead - del model - - for m in self.model_parts: - m.to_empty(device=init_device) - with torch.no_grad(): - m.init_states(buffer_device=buffer_device) - m.train() - else: - # apply PT-D Tensor Parallel, activation checkpointing, torch.compile, Data Parallel - model = self.train_spec.parallelize_fn( - model, - parallel_dims=parallel_dims, - training=config.training, - parallelism=config.parallelism, - compile_config=config.compile, - ac_config=config.activation_checkpoint, - dump_folder=config.dump_folder, - ) - - model.to_empty(device=init_device) - with torch.no_grad(): - model.init_states(buffer_device=buffer_device) - model.train() - - self.model_parts = [model] - - # build optimizer after applying parallelisms to the model - self.optimizers = config.optimizer.build( - model_parts=self.model_parts, - ) - if self.train_spec.post_optimizer_build_fn is not None: - self.train_spec.post_optimizer_build_fn( - self.optimizers, self.model_parts, parallel_dims - ) - self.lr_schedulers = config.lr_scheduler.build( - optimizers=self.optimizers, - training_steps=config.training.steps, - ) - - self.checkpointer = config.checkpoint.build( - dataloader=None, - model_parts=self.model_parts, - optimizers=self.optimizers, - lr_schedulers=self.lr_schedulers, - states={"train_state": self}, - sd_adapter=( - self.train_spec.state_dict_adapter(model_config, config.hf_assets_path) - if self.train_spec.state_dict_adapter - else None - ), - base_folder=config.dump_folder, - ) - - self.train_context = dist_utils.get_spmd_context( - parallel_dims=parallel_dims, - ) - - def close(self) -> None: - if self.checkpointer: - self.checkpointer.close() diff --git a/torchtitan/experiments/forge/example_train.py b/torchtitan/experiments/forge/example_train.py deleted file mode 100644 index 09b9e2b39a..0000000000 --- a/torchtitan/experiments/forge/example_train.py +++ /dev/null @@ -1,469 +0,0 @@ -# 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 os -import time -from collections.abc import Iterable -from datetime import timedelta -from typing import Any - -import torch -from torch.distributed.elastic.multiprocessing.errors import record - -from torchtitan.components.data.loader import BaseDataLoader, DataloaderExhaustedError -from torchtitan.components.metrics import MetricsProcessor -from torchtitan.components.tokenizer import HuggingFaceTokenizer -from torchtitan.components.validate import Validator -from torchtitan.config import ConfigManager -from torchtitan.distributed import utils as dist_utils -from torchtitan.distributed.context_parallel import prepare_context_parallel_input -from torchtitan.tools import utils -from torchtitan.tools.logging import init_logger, logger -from torchtitan.trainer import Trainer as TitanTrainer - -from .engine import ForgeEngine - - -class Trainer(ForgeEngine): - tokenizer: HuggingFaceTokenizer | None - dataloader: BaseDataLoader - validator: Validator - metrics_processor: MetricsProcessor - - # additional training states - step: int - - # Enable debug tracing on failure: https://pytorch.org/docs/stable/elastic/errors.html - @record - def __init__(self, config: TitanTrainer.Config): - if config.debug.print_config: - logger.info(f"Running with args: {config.to_dict()}") - - # NOTE: Here we are passing in Trainer.Config as a superset of ForgeEngine.Config - super().__init__(config) - - # build tokenizer - self.tokenizer = ( - config.tokenizer.build(tokenizer_path=config.hf_assets_path) - if config.tokenizer is not None - else None - ) - - # build dataloader - num_tokens_per_batch = config.training.num_tokens_per_microbatch_per_dp_rank - self.dataloader = config.dataloader.build( - dp_world_size=self.dp_degree, - dp_rank=self.dp_rank, - tokenizer=self.tokenizer, - max_context_length=config.training.max_context_length, - num_tokens_per_batch=num_tokens_per_batch, - ) - - model_args = self.model_config - logger.info( - f"Built {config.model_spec.name} {config.model_spec.flavor} with {model_args}" - ) - - # metrics logging - self.metrics_processor = config.metrics.build( - parallel_dims=self.parallel_dims, - dump_folder=config.dump_folder, - pp_schedule=config.parallelism.pipeline_parallel_schedule, - config_dict=config.to_dict(), - ) - color = self.metrics_processor.color - - self.metrics_processor.num_flops_per_token = self.num_flops_per_token - - logger.info( - f"{color.blue}Model {config.model_spec.name} {config.model_spec.flavor} " - f"{color.red}size: {self.model_param_count:,} total parameters{color.reset}" - ) - - # initialize device memory monitor and get peak flops for MFU calculation - device_memory_monitor = self.metrics_processor.device_memory_monitor - gpu_peak_flops = utils.get_peak_flops(device_memory_monitor.device_name) - logger.info(f"Peak FLOPS used for computing MFU: {gpu_peak_flops:.3e}") - device_mem_stats = device_memory_monitor.get_peak_stats() - logger.info( - f"{utils.device_type.upper()} memory usage for model: " - f"{device_mem_stats.max_reserved_gib:.2f}GiB" - f"({device_mem_stats.max_reserved_pct:.2f}%)" - ) - - self.metrics_processor.optimizers = self.optimizers - - # Initialize trainer states that will be saved in checkpoint. - # These attributes must be initialized before checkpoint loading. - self.step = 0 - - # Build validator if validation is configured - if config.validator.enable: - pp_schedule, pp_has_first_stage, pp_has_last_stage = ( - ( - self.pp_schedule, - self.pp_has_first_stage, - self.pp_has_last_stage, - ) - if self.parallel_dims.pp_enabled - else (None, None, None) - ) - - self.validator = config.validator.build( - parallelism=config.parallelism, - dp_world_size=self.dp_degree, - dp_rank=self.dp_rank, - tokenizer=self.tokenizer, - parallel_dims=self.parallel_dims, - loss_fn=self.loss_fn, - validation_context=self.train_context, - metrics_processor=self.metrics_processor, - seq_len=config.training.max_context_length, - num_tokens_per_batch=num_tokens_per_batch, - pp_schedule=pp_schedule, - pp_has_first_stage=pp_has_first_stage, - pp_has_last_stage=pp_has_last_stage, - ) - - self.profiler = config.profiler.build() - - logger.info( - "Trainer is initialized with " - f"{config.training.num_tokens_per_microbatch_per_dp_rank * self.num_pp_microbatches} " - "tokens per DP rank, " - f"{self.num_tokens_per_train_step} tokens per train step, " - f"gradient accumulation steps {self.gradient_accumulation_steps}, " - f"maximum context length {config.training.max_context_length}, " - f"total steps {config.training.steps} " - f"(warmup {config.lr_scheduler.warmup_steps})." - ) - - def batch_generator( - self, data_iterable: Iterable[tuple[dict[str, torch.Tensor], torch.Tensor]] - ) -> Iterable[tuple[dict[str, torch.Tensor], torch.Tensor]]: - """Returns an iterator that processes batches from the data iterator.""" - device_type = utils.device_type - data_iterator = iter(data_iterable) - - while True: - try: - batch = next(data_iterator) - except StopIteration as ex: - # If data runs out during gradient accumulation, that - # entire step will not be executed. - raise DataloaderExhaustedError() from ex - data_load_start = time.perf_counter() - input_dict, labels = batch - self.metrics_processor.ntokens_since_last_log += labels.numel() - self.metrics_processor.data_loading_times.append( - time.perf_counter() - data_load_start - ) - - # Tensors stay on CPU; moved to GPU per-microbatch during training - yield input_dict, labels - - def post_dataloading_process( - self, input_dict: dict[str, torch.Tensor], labels: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: - inputs = input_dict["input"] - # Everything except the pipelined input is a model-forward kwarg, - # forwarded to all PP stages by the schedule. - extra_kwargs: dict[str, Any] = { - k: v for k, v in input_dict.items() if k != "input" - } - - positions = extra_kwargs.get("positions", None) - - try: - # pyrefly: ignore [not-callable] - extra_kwargs["attention_masks"] = self.model_parts[0].get_attention_masks( - positions=positions, - ) - except TypeError: - pass - - if self.parallel_dims.cp_enabled: - cp_input_dict = prepare_context_parallel_input( - {"input": inputs, "labels": labels, **extra_kwargs}, - None, - self.parallel_dims.get_mesh("cp"), - self.config.parallelism.context_parallel_load_balancer, - self.config.parallelism.context_parallel_ptrr_mask_key, - ) - inputs = cp_input_dict.pop("input") - labels = cp_input_dict.pop("labels") - extra_kwargs = cp_input_dict - - return inputs, labels, extra_kwargs - - def forward_backward_step( - self, - *, - input_dict: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]], - labels: torch.Tensor | list[torch.Tensor], - global_valid_tokens: torch.Tensor, - ) -> torch.Tensor: - model_parts = self.model_parts - parallel_dims = self.parallel_dims - - if parallel_dims.pp_enabled: - assert isinstance(input_dict, list) - assert isinstance(labels, list) - return self.pp_forward_backward_step( - input_dict_mbs=input_dict, - label_mbs=labels, - global_valid_tokens=global_valid_tokens, - ) - - assert isinstance(input_dict, dict) - assert isinstance(labels, torch.Tensor) - inputs, labels, extra_kwargs = self.post_dataloading_process(input_dict, labels) - - with self.train_context(): - assert len(model_parts) == 1 - pred = model_parts[0](inputs, **extra_kwargs) - loss_sum, _ = self.loss_fn(pred, labels) - loss = loss_sum / global_valid_tokens - del pred - loss.backward() - - return loss - - def pp_forward_backward_step( - self, - *, - input_dict_mbs: list[dict[str, torch.Tensor]], - label_mbs: list[torch.Tensor], - global_valid_tokens: torch.Tensor, - ) -> torch.Tensor: - arg_mbs: list[tuple[torch.Tensor, ...]] = [] - kwarg_mbs: list[dict[str, Any]] = [] - target_mbs: list[torch.Tensor] | None = [] if self.pp_has_last_stage else None - for input_dict, labels in zip(input_dict_mbs, label_mbs, strict=True): - inputs, labels, extra_kwargs = self.post_dataloading_process( - input_dict, labels - ) - if self.pp_has_first_stage: - arg_mbs.append((inputs,)) - kwarg_mbs.append(extra_kwargs) - if target_mbs is not None: - target_mbs.append(labels) - - loss_kwargs = {"global_valid_tokens": global_valid_tokens} - with self.train_context(): - losses = [] if self.pp_has_last_stage else None - self.pp_schedule.step( - arg_mbs=arg_mbs if self.pp_has_first_stage else None, - kwarg_mbs=kwarg_mbs, - target_mbs=target_mbs, - losses=losses, - loss_kwargs=loss_kwargs, - return_outputs=False, - ) - - # TODO: PP+FSDP unexpectedly puts the loss back to the CPU. - if self.pp_has_last_stage: - assert losses is not None - return torch.sum(torch.stack(losses)).to(self.device) - return torch.tensor([-1.0], device=self.device) - - def train_step( - self, data_iterator: Iterable[tuple[dict[str, torch.Tensor], torch.Tensor]] - ): - self.optimizers.zero_grad() - - # Keep these variables local to shorten the code as these are - # the major variables that are used in the training loop. - parallel_dims = self.parallel_dims - # All groups form one optimizer step; each group feeds one fwd-bwd call. - microbatch_groups: list[list[tuple[dict[str, torch.Tensor], torch.Tensor]]] = [] - local_valid_tokens = torch.tensor(0, dtype=torch.int64) - for _ in range(self.gradient_accumulation_steps): - microbatches = [] - for _ in range(self.num_pp_microbatches): - input_dict, labels = next(data_iterator) - # Popped so the batch reaching the model holds only its kwargs. - local_valid_tokens += input_dict.pop("num_valid_tokens") - microbatches.append((input_dict, labels)) - microbatch_groups.append(microbatches) - - # Keep the global token count on device so loss normalization does not - # introduce a CPU synchronization in the training path. - global_valid_tokens = local_valid_tokens.to(self.device) - if parallel_dims.dp_enabled: - batch_mesh = parallel_dims.get_mesh("batch") - global_valid_tokens = dist_utils.dist_sum_tensor( - global_valid_tokens, batch_mesh - ) - - accumulated_losses = [] - for microbatches in microbatch_groups: - input_dict_mbs = [] - label_mbs = [] - for input_dict, labels in microbatches: - for key, value in input_dict.items(): - if isinstance(value, torch.Tensor): - input_dict[key] = value.to(self.device) - input_dict_mbs.append(input_dict) - label_mbs.append(labels.to(self.device)) - - if parallel_dims.pp_enabled: - fwd_bwd_input_dict = input_dict_mbs - fwd_bwd_labels = label_mbs - else: - assert len(input_dict_mbs) == len(label_mbs) == 1 - fwd_bwd_input_dict = input_dict_mbs[0] - fwd_bwd_labels = label_mbs[0] - - loss = self.forward_backward_step( - input_dict=fwd_bwd_input_dict, - labels=fwd_bwd_labels, - global_valid_tokens=global_valid_tokens, - ) - accumulated_losses.append(loss.detach()) - - grad_norm = dist_utils.clip_grad_norm_( - [p for m in self.model_parts for p in m.parameters()], - self.config.training.max_norm, - foreach=True, - pp_mesh=parallel_dims.get_optional_mesh("pp"), - ep_enabled=parallel_dims.ep_enabled, - ) - self.checkpointer.maybe_wait_for_staging() - self.optimizers.step() - self.lr_schedulers.step() - - # Reduce the data collected over gradient accumulation steps. - loss = torch.sum(torch.stack(accumulated_losses)) - - # log metrics - if not self.metrics_processor.should_log(self.step): - return - - if parallel_dims.dp_cp_enabled: - loss = loss.detach() - global_avg_loss, global_max_loss = ( - dist_utils.dist_sum(loss, parallel_dims.get_optional_mesh("loss")), - dist_utils.dist_max(loss, parallel_dims.get_optional_mesh("loss")), - ) - else: - global_avg_loss = global_max_loss = loss.detach().item() - - self.metrics_processor.log( - self.step, - global_avg_loss, - global_max_loss, - grad_norm.item(), - ) - - @record - def train(self): - config = self.config - - self.checkpointer.load(step=config.checkpoint.load_step) - logger.info(f"Training starts at step {self.step + 1}.") - - with self.profiler.active( - global_step=self.step, - base_folder=config.dump_folder, - ) as profiler: - data_iterator = self.batch_generator(self.dataloader) - while self.step < config.training.steps: - self.step += 1 - self.gc_handler.run(self.step) - try: - self.train_step(data_iterator) - except DataloaderExhaustedError: - logger.warning("Ran out of data; last step was canceled.") - break - - # Run validation if validator is available - if config.validator.enable and self.validator.should_validate( - self.step - ): - self.validator.validate(self.model_parts, self.step) - - self.checkpointer.save( - self.step, last_step=(self.step == config.training.steps) - ) - - # signal the profiler that the next profiling step has started - profiler.step() - - # reduce timeout after first train step for faster signal - # (assuming lazy init and compilation are finished) - if self.step == 1: - dist_utils.set_pg_timeouts( - timeout=timedelta(seconds=config.comm.train_timeout_seconds), - parallel_dims=self.parallel_dims, - ) - - if torch.distributed.get_rank() == 0: - logger.info("Sleeping 2 seconds for other ranks to complete") - time.sleep(2) - - logger.info("Training completed") - - def state_dict(self) -> dict[str, Any]: - return {"step": self.step} - - def load_state_dict(self, state_dict: dict[str, Any]): - self.step = state_dict["step"] - - def close(self) -> None: - if self.metrics_processor: - self.metrics_processor.close() - super().close() - - -def main(custom_trainer_class: type[Trainer] | None = None) -> None: - """Main entry point for training.""" - init_logger() - - import torchtitan - - logger.info( - "torchtitan version: %s (0.0.0 means __version__ is not defined correctly).", - torchtitan.__version__, - ) - - config_manager = ConfigManager() - config = config_manager.parse_args() - trainer: Trainer | None = None - - try: - # pyrefly: ignore [missing-attribute] - if custom_trainer_class is not None: - trainer = custom_trainer_class(config) - else: - trainer = config.build() - - # pyrefly: ignore [missing-attribute] - if config.checkpoint.create_seed_checkpoint: - assert ( - int(os.environ["WORLD_SIZE"]) == 1 - ), "Must create seed checkpoint using a single device, to disable sharding." - assert ( - # pyrefly: ignore [missing-attribute] - config.checkpoint.enable - ), "Must enable checkpointing when creating a seed checkpoint." - trainer.checkpointer.save(curr_step=0, last_step=True) - logger.info("Created seed checkpoint") - else: - trainer.train() - except Exception: - if trainer: - trainer.close() - raise - else: - trainer.close() - if torch.distributed.is_initialized(): - torch.distributed.destroy_process_group() - logger.info("Process group destroyed") - - -if __name__ == "__main__": - main(Trainer) diff --git a/torchtitan/experiments/graph_trainer/.claude/CLAUDE.md b/torchtitan/experiments/graph_trainer/.claude/CLAUDE.md index dcb0b1034d..7b03f26650 100644 --- a/torchtitan/experiments/graph_trainer/.claude/CLAUDE.md +++ b/torchtitan/experiments/graph_trainer/.claude/CLAUDE.md @@ -340,10 +340,9 @@ NGPU=4 MODULE=graph_trainer.llama3 CONFIG=graph_trainer_llama3_8b_c4_test ./run_ The `insert_kernel_annotations_pass` labels CUDA graph kernels with their originating `nn.Module` path in profiler traces. It runs automatically in the -`aot_fx_trace` path (bundled with the cudagraph pass). The profiler always runs -``cudagraph_annotate_trace_post_processor`` after exporting a trace, so CUDA -graph annotations are merged automatically and no manual post-processing is -needed. +`aot_fx_trace` path (bundled with the cudagraph pass). The profiler passes the +captured annotations to ``export_chrome_trace``, which bakes them into the trace +as it writes, so they are merged automatically and no post-processing is needed. Requirements: `cuda-python` package and CUDA toolkit/driver >= 13.1 (or `cuda-compat >= 13.1` on `LD_LIBRARY_PATH`). The pass is a no-op when diff --git a/torchtitan/experiments/graph_trainer/common_utils.py b/torchtitan/experiments/graph_trainer/common_utils.py index 312fb84b17..b89903959f 100644 --- a/torchtitan/experiments/graph_trainer/common_utils.py +++ b/torchtitan/experiments/graph_trainer/common_utils.py @@ -35,6 +35,22 @@ AnnotatedLossFn: TypeAlias = Callable[..., LossResult] +def _get_graph_modules( + gm: torch.fx.GraphModule, + *, + recurse: bool, + apply_to_root: bool, +) -> list[torch.fx.GraphModule]: + modules = [gm] if apply_to_root else [] + if recurse: + modules.extend( + module + for name, module in gm.named_modules() + if name and isinstance(module, torch.fx.GraphModule) + ) + return modules + + class GraphTrainerScaledDotProductAttention(ScaledDotProductAttention): """Adapt flat graph-trainer attention inputs to the batched SDPA interface.""" diff --git a/torchtitan/experiments/graph_trainer/configs.py b/torchtitan/experiments/graph_trainer/configs.py index 71941f2ae1..8ab1da3e5a 100644 --- a/torchtitan/experiments/graph_trainer/configs.py +++ b/torchtitan/experiments/graph_trainer/configs.py @@ -94,7 +94,9 @@ class GraphTrainerCompileConfig(CompileConfig): debug_graph_passes: bool = False """Log timing, op-count diffs, and before/after graphs for each pass to tlparse.""" - memory_policy: Literal["default", "full", "eager", "sac_and_offload"] = "default" + memory_policy: Literal[ + "default", "full", "eager", "min_cut", "sac_and_offload" + ] = "default" """ Memory optimization policy for activation management (SAC, offload). default: SAC — save all compute-intensive ops and FSDP all_gathers. @@ -103,6 +105,7 @@ class GraphTrainerCompileConfig(CompileConfig): full AC (checkpoint_wrapper with no context_fn). eager: SAC alternating mm ops between save/recompute, matching the eager AC policy in torchtitan.distributed.activation_checkpoint. + min_cut: choose saved activations with the min-cut partitioner. sac_and_offload: SAC + CPU offload — apply default SAC first, then offload surviving MUST_SAVE activations to CPU within the cpu_offload_budget_gb budget. diff --git a/torchtitan/experiments/graph_trainer/decompositions.py b/torchtitan/experiments/graph_trainer/decompositions.py new file mode 100644 index 0000000000..05a7be476f --- /dev/null +++ b/torchtitan/experiments/graph_trainer/decompositions.py @@ -0,0 +1,100 @@ +# 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. + +"""Standalone decomposition passes for GraphTrainer FX graphs.""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import Any + +import torch +import torch.fx as fx +import torch.fx.traceback as fx_traceback +from torch.fx.experimental.proxy_tensor import selective_decompose + +from torchtitan.experiments.graph_trainer.common_utils import _get_graph_modules + + +def _copy_attr(src: torch.nn.Module, dst: torch.nn.Module, target: str) -> None: + *prefix, field = target.split(".") + for name in prefix: + src_child = getattr(src, name) + child = getattr(dst, name, None) + if src_child is child: + return + if child is None: + child = torch.nn.Module() + setattr(dst, name, child) + src = src_child + dst = child + + value = getattr(src, field) + if isinstance(value, torch.Tensor) and not isinstance(value, torch.nn.Parameter): + persistent = field not in src._non_persistent_buffers_set + dst.register_buffer(field, value, persistent=persistent) + else: + setattr(dst, field, value) + + +def _apply_decompositions( + gm: fx.GraphModule, + example_inputs, + decomposition_table: dict[torch._ops.OperatorBase, Callable[..., Any]], +) -> None: + if example_inputs is None: + placeholders = gm.graph.find_nodes(op="placeholder") + if any("val" not in node.meta for node in placeholders): + return + example_inputs = tuple(node.meta["val"] for node in placeholders) + + with fx_traceback.preserve_node_meta(): + decomposed = selective_decompose( + gm, + *example_inputs, + decomposition=decomposition_table, + should_decompose=lambda _node: True, + trace_joint_graph=False, + ) + + for node in decomposed.graph.find_nodes(op="get_attr"): + _copy_attr(decomposed, gm, node.target) + gm.graph = decomposed.graph + gm.graph.lint() + gm.recompile() + + +def apply_decompositions_pass( + gm: fx.GraphModule, + example_inputs=None, + *, + decomposition_table: dict[torch._ops.OperatorBase, Callable[..., Any]], + recurse: bool = False, + apply_to_root: bool = True, +) -> fx.GraphModule: + """Apply decompositions to selected FX graph modules in place. + + Args: + gm: Root graph module. + example_inputs: Inputs for the root graph. If omitted, inputs are read + from placeholder metadata. + decomposition_table: Operator-to-decomposition mapping. + recurse: Whether to apply decompositions to nested graph modules. + apply_to_root: Whether to apply decompositions to ``gm`` itself. + + Returns: + The input graph module, modified in place. + """ + if not decomposition_table: + return gm + + for module in _get_graph_modules(gm, recurse=recurse, apply_to_root=apply_to_root): + _apply_decompositions( + module, + example_inputs if module is gm else None, + decomposition_table, + ) + return gm diff --git a/torchtitan/experiments/graph_trainer/deepseek_v3/config_registry.py b/torchtitan/experiments/graph_trainer/deepseek_v3/config_registry.py index e6b14f54cb..291536945b 100644 --- a/torchtitan/experiments/graph_trainer/deepseek_v3/config_registry.py +++ b/torchtitan/experiments/graph_trainer/deepseek_v3/config_registry.py @@ -6,10 +6,7 @@ from dataclasses import replace -from torchtitan.components.quantization import ( - MXFP8GroupedExpertsConverter, - MXFP8LinearConverter, -) +from torchtitan.components.quantization import MXFP8GroupedExpertsConverter from torchtitan.distributed.pipeline_parallel import pipeline_llm from torchtitan.experiments.graph_trainer.configs import ( GraphTrainerCompileConfig, @@ -23,6 +20,7 @@ deepseek_v3_671b, deepseek_v3_debugmodel, deepseek_v3_debugmodel_minimal_async_ep, + deepseek_v3_mxfp8_linear_converter_config, ) from . import model_registry @@ -42,9 +40,8 @@ def graph_trainer_deepseek_v3_debugmodel_mxfp8() -> GraphTrainer.Config: base.model_spec = deepseek_v3_model_registry( "debugmodel", converters=[ - MXFP8LinearConverter.Config( + deepseek_v3_mxfp8_linear_converter_config( model_compile_enabled=True, - fqns=["attention", "shared_experts", "feed_forward"], ), MXFP8GroupedExpertsConverter.Config( model_compile_enabled=True, diff --git a/torchtitan/experiments/graph_trainer/llama3/config_registry.py b/torchtitan/experiments/graph_trainer/llama3/config_registry.py index d2c5dcbe45..0ccbd0804b 100644 --- a/torchtitan/experiments/graph_trainer/llama3/config_registry.py +++ b/torchtitan/experiments/graph_trainer/llama3/config_registry.py @@ -8,7 +8,6 @@ from torchtitan.components.data import ConcatThenSplitPackingConfig, GrainDataLoader from torchtitan.components.loss import CrossEntropyLoss -from torchtitan.components.quantization import MXFP8LinearConverter from torchtitan.experiments.graph_trainer.configs import ( GraphTrainerCompileConfig, to_graph_trainer_config, @@ -23,6 +22,7 @@ llama3_8b, llama3_debugmodel, llama3_debugmodel_dist_gemm, + llama3_mxfp8_linear_converter_config, ) from torchtitan.observability.sdc_replayer import SDCReplayer @@ -61,6 +61,19 @@ def graph_trainer_llama3_debugmodel_dist_gemm() -> GraphTrainer.Config: return config +def graph_trainer_llama3_debugmodel_mxfp8() -> GraphTrainer.Config: + base = llama3_debugmodel() + base.model_spec = llama3_model_registry( + "debugmodel", + converters=[ + llama3_mxfp8_linear_converter_config(model_compile_enabled=True), + ], + ) + config = to_graph_trainer_config(base, model_registry) + config.compile = GraphTrainerCompileConfig(enable=True) + return config + + def graph_trainer_llama3_debugmodel_sdpa() -> GraphTrainer.Config: """Debug model on the test-only SDPA backend. @@ -127,7 +140,9 @@ def graph_trainer_llama3_8b_mxfp8() -> GraphTrainer.Config: # MXFP8 converter's compile requirement is satisfied. base.model_spec = llama3_model_registry( "8B", - converters=[MXFP8LinearConverter.Config(model_compile_enabled=True)], + converters=[ + llama3_mxfp8_linear_converter_config(model_compile_enabled=True), + ], ) config = to_graph_trainer_config(base, model_registry) config.compile = GraphTrainerCompileConfig(enable=True) diff --git a/torchtitan/experiments/graph_trainer/memory_policy.py b/torchtitan/experiments/graph_trainer/memory_policy.py index ba26f1d29d..cf7cba0ee2 100644 --- a/torchtitan/experiments/graph_trainer/memory_policy.py +++ b/torchtitan/experiments/graph_trainer/memory_policy.py @@ -21,6 +21,15 @@ from typing import TYPE_CHECKING import torch +from torch._functorch.partitioners import ( + choose_saved_values_set, + force_save_bw_mutation_src, + force_save_collectives, + force_save_effectful_ops, + get_default_op_list, + NodeInfo, +) +from torch.utils._ordered_set import OrderedSet from torch.utils.checkpoint import CheckpointPolicy from torchtitan.distributed.activation_checkpoint import _get_default_save_ops @@ -52,6 +61,9 @@ from torchtitan.experiments.graph_trainer.configs import GraphTrainerCompileConfig +_INF_DISTANCE = int(1e9) + + def _make_default_memory_policy(save_ops: set | None = None) -> Callable: """Create a SAC policy function from a set of op targets to save.""" if save_ops is None: @@ -389,6 +401,144 @@ def _eager_memory_policy_pass( return gm +def _is_backward_side(node: torch.fx.Node, backward_side: set[torch.fx.Node]) -> bool: + return _is_backward_node(node) or any( + inp in backward_side for inp in node.all_input_nodes + ) + + +def _backward_side_nodes( + gm: torch.fx.GraphModule, +) -> OrderedSet[torch.fx.Node]: + backward_side = OrderedSet() + for node in gm.graph.nodes: + if node.op != "output" and _is_backward_side(node, backward_side): + backward_side.add(node) + return backward_side + + +def _node_info_for_graph_trainer( + gm: torch.fx.GraphModule, + backward_side: OrderedSet[torch.fx.Node], +) -> NodeInfo | None: + nodes = list(gm.graph.nodes) + required_bw_nodes = OrderedSet( + node for node in nodes if node in backward_side and node.op != "output" + ) + if not required_bw_nodes: + return None + + required_fw_nodes = OrderedSet( + node for node in nodes if node not in required_bw_nodes and node.op != "output" + ) + fw_order = {node: idx for idx, node in enumerate(required_fw_nodes)} + static_lifetime_input_nodes = OrderedSet( + node for node in required_fw_nodes if node.op in ("placeholder", "get_attr") + ) + + for node in reversed(nodes): + if node.op == "output": + node.dist_from_bw = _INF_DISTANCE + elif node in required_bw_nodes: + node.dist_from_bw = 0 + elif node in required_fw_nodes: + user_distances = [ + getattr(user, "dist_from_bw", _INF_DISTANCE) + 1 for user in node.users + ] + node.dist_from_bw = min(user_distances, default=_INF_DISTANCE) + else: + node.dist_from_bw = _INF_DISTANCE + + return NodeInfo( + list(static_lifetime_input_nodes), + required_fw_nodes, + required_bw_nodes.copy(), + required_bw_nodes.copy(), + OrderedSet(), + fw_order, + static_lifetime_input_nodes, + ) + + +def tag_min_cut_saved_values( + gm: torch.fx.GraphModule, + backward_side: OrderedSet[torch.fx.Node], + saved_values: set[torch.fx.Node], +) -> None: + required_fw_nodes = { + node + for node in gm.graph.nodes + if node not in backward_side and node.op != "output" + } + saved_boundaries = set(saved_values) + op_types = get_default_op_list() + pending = list(saved_boundaries) + while pending: + node = pending.pop() + if node not in saved_boundaries: + continue + if node not in required_fw_nodes or node.op != "call_function": + continue + if ( + node.target == torch.ops.aten.detach.default or op_types.is_view(node) + ) and any(inp in required_fw_nodes for inp in node.all_input_nodes): + saved_boundaries.remove(node) + for inp in node.all_input_nodes: + if inp in required_fw_nodes and inp not in saved_boundaries: + saved_boundaries.add(inp) + pending.append(inp) + + saved_boundaries.update( + node + for node in required_fw_nodes + if node.meta.get("recompute") == CheckpointPolicy.MUST_SAVE + ) + for node in saved_boundaries: + if node in required_fw_nodes: + node.meta["recompute"] = CheckpointPolicy.MUST_SAVE + + seen = set() + + def visit(node: torch.fx.Node) -> None: + if node in seen or node in saved_boundaries: + return + seen.add(node) + if node in backward_side: + for inp in node.all_input_nodes: + visit(inp) + return + if node not in required_fw_nodes or node.op in ("placeholder", "get_attr"): + return + if node.op == "call_function": + node.meta["recompute"] = CheckpointPolicy.MUST_RECOMPUTE + for inp in node.all_input_nodes: + visit(inp) + + for node in backward_side: + for inp in node.all_input_nodes: + visit(inp) + + +@register_memory_policy("min_cut") +def _min_cut_memory_policy_pass( + gm: torch.fx.GraphModule, + *, + config: "GraphTrainer.Config", +) -> torch.fx.GraphModule: + """Choose saved activations with the min-cut partitioner.""" + backward_side = _backward_side_nodes(gm) + node_info = _node_info_for_graph_trainer(gm, backward_side) + if node_info is None: + return gm + + force_save_collectives(gm) + force_save_effectful_ops(gm) + force_save_bw_mutation_src(gm) + saved_values = choose_saved_values_set(gm.graph, node_info) + tag_min_cut_saved_values(gm, backward_side, set(saved_values)) + return gm + + @register_memory_policy("sac_and_offload") def _sac_and_offload_memory_policy_pass( gm: torch.fx.GraphModule, @@ -416,10 +566,12 @@ def tag_with_memory_policy_pass( default: SAC with all compute-intensive ops saved. full: full recompute except user-selected module operations. eager: SAC alternating mm ops between save/recompute. + min_cut: choose saved activations with the min-cut partitioner. sac_and_offload: SAC + CPU offload within budget. Other memory policies combining SAC and CPU offload can be added via ``register_memory_policy`` without modifying this function. + """ memory_policy = config.compile.memory_policy if memory_policy not in MEMORY_POLICY_REGISTRY: @@ -427,6 +579,7 @@ def tag_with_memory_policy_pass( f"Unknown memory_policy: {memory_policy!r}. " f"Available: {list(MEMORY_POLICY_REGISTRY.keys())}" ) + gm = MEMORY_POLICY_REGISTRY[memory_policy](gm, config=config) log_activation_memory_policy(gm) return gm diff --git a/torchtitan/experiments/graph_trainer/passes.py b/torchtitan/experiments/graph_trainer/passes.py index 5a8c2291eb..16a087855c 100644 --- a/torchtitan/experiments/graph_trainer/passes.py +++ b/torchtitan/experiments/graph_trainer/passes.py @@ -10,7 +10,8 @@ This module provides pass orchestration: building the pass list, applying passes in order, and the pass registries. Individual passes live in dedicated modules: -- ``memory_policy.py`` — SAC tagging and memory policy dispatch +- ``memory_policy.py`` - SAC and min-cut policy tagging and dispatch +- ``decompositions.py`` — standalone graph decomposition - ``inductor_passes.py`` — regional and full Inductor compilation - ``cudagraph.py`` — cudagraph wrapping and kernel annotations - ``fsdp_passes.py`` — FSDP bucketing and resharding @@ -18,6 +19,8 @@ cleanup bundled as ``canonicalize_graph_pass`` (detach, identity view/slice, back-to-back transpose, view→reshape normalization) - ``performance_passes.py`` — opt-in numerics-changing optimizations +- ``subgraph_regions.py`` — region annotation, invoke_subgraph outlining, and + shared region prologue extraction - ``selective_activation_remat.py`` — activation rematerialization - ``cpu_offload.py`` — CPU offload insertion - ``custom_codegen.py`` — custom code generation for profiling/debugging diff --git a/torchtitan/experiments/graph_trainer/simple_fsdp.py b/torchtitan/experiments/graph_trainer/simple_fsdp.py index c36176cfdd..538dd1ca60 100644 --- a/torchtitan/experiments/graph_trainer/simple_fsdp.py +++ b/torchtitan/experiments/graph_trainer/simple_fsdp.py @@ -27,6 +27,10 @@ from torch.distributed.tensor._redistribute import redistribute_local_tensor from torch.distributed.tensor.placement_types import _StridedShard, Placement +from torchtitan.components.quantization._fsdp_tensor import ( + _ShardedFSDPTensor, + _UnshardedFSDPTensor, +) from torchtitan.distributed.utils import get_spmd_backend from torchtitan.protocols.module import Module @@ -222,6 +226,46 @@ def _register_parametrization( module.__class__ = module_cls +class _BuildUnshardedTensorFunction(torch.autograd.Function): + """Own the gradient edge for an unsharded tensor built outside FSDP2. + + FSDP2 creates this edge internally for its post-all-gather output. + GraphTrainer's SimpleFSDP reconstructs the unsharded weight itself and so + has no such edge, and this routes its logical tensor gradient straight + back through the high-precision gather. + + Lives here, with its only user, rather than beside the FSDP2 edge it + mirrors. ``forward`` does the same two steps ``fsdp_post_all_gather`` does + on its first unshard -- quantize, then wrap. + """ + + @staticmethod + # pyrefly: ignore [bad-override] + def forward(ctx, weight: torch.Tensor, wrapper: _ShardedFSDPTensor): + del ctx + # Unreachable today, so this is an invariant rather than a fallback: + # replicate_compute returns a plain local tensor under spmd_types, only + # partial_dtensor composed with TP or EP re-wraps it on the + # non-data-parallel mesh, and MXFP8Linear.Config.build rejects that + # backend outright. The two guards sit in different files, so assert + # here: a future format subclassing _ShardedFSDPTensor without that + # rejection should fail loudly rather than silently quantize a + # DTensor's local shard. + assert not isinstance(weight, DTensor), ( + "unsharded tensor received a DTensor, so a format on this " + "lifecycle reached the partial_dtensor backend without rejecting " + "it; see MXFP8Linear.Config.build" + ) + with torch.no_grad(): + return _UnshardedFSDPTensor(weight, wrapper._build_operands(weight)) + + @staticmethod + # pyrefly: ignore [bad-override] + def backward(ctx, grad_weight: torch.Tensor): + del ctx + return grad_weight, None + + class ReplicateComputation(Module): def __init__( self, @@ -328,8 +372,31 @@ def forward(self, x: DTensor) -> torch.Tensor: if not _active_parametrization: return x - output = self.replicate_compute(x) - return output + unsharded_weight = self.replicate_compute(x) + # Which operands to build is determined by the tensor subclass on the + # *sharded* parameter, and replicate_compute does not preserve it: its + # return value is a plain local tensor under spmd_types. Its input is + # not -- ``x`` is the sharded parameter, still a DTensor, because + # data_parallel translates parameters to full-mesh DTensors before + # this runs (replicate_compute reads ``x._spec`` on entry for the same + # reason). So take the subclass from ``x`` and the values from the + # return value. A parameter with no operands passes straight through, + # so this costs nothing when nothing is quantized. + # + # Read the local tensor directly rather than through ``to_local``: + # only the subclass type is wanted, so the autograd edge and the extra + # traced node that ``to_local`` adds would both be dead weight. + source = x._local_tensor + if isinstance(source, _UnshardedFSDPTensor): + raise RuntimeError( + "The data parallel parametrization received an already-" + "unsharded weight. FSDP2 builds the unsharded tensor in " + "fsdp_post_all_gather and owns its gradient edge, so building " + "one here as well would quantize the weight a second time." + ) + if not isinstance(source, _ShardedFSDPTensor): + return unsharded_weight + return _BuildUnshardedTensorFunction.apply(unsharded_weight, source) def data_parallel( @@ -343,6 +410,7 @@ def data_parallel( # TODO: Unify this with device_mesh as a global data- and model-parallel mesh. non_dp_mesh: DeviceMesh | None = None, ) -> nn.Module: + """Shard ``model`` and install the data-parallel parametrization.""" param_sharding: tuple[Placement, ...] if mode == "replicate": param_sharding = (Replicate(),) diff --git a/torchtitan/experiments/graph_trainer/subgraph_regions.py b/torchtitan/experiments/graph_trainer/subgraph_regions.py new file mode 100644 index 0000000000..004fa18c07 --- /dev/null +++ b/torchtitan/experiments/graph_trainer/subgraph_regions.py @@ -0,0 +1,622 @@ +# 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. + +from __future__ import annotations + +from contextlib import nullcontext +from operator import attrgetter, getitem +from typing import Any + +import torch +import torch.fx as fx +from torch.fx import Node +from torch.fx._lazy_graph_module import _LazyGraphModule +from torch.fx.traceback import annotate +from torch.utils._ordered_set import OrderedSet + +SUBGRAPH_REGION = "graph_trainer_subgraph_region" +SUBGRAPH_REGION_ROLE = "graph_trainer_subgraph_region_role" +SUBGRAPH_REGION_PRESERVE_ORDER = "graph_trainer_subgraph_region_preserve_order" +_SUBGRAPH_REGION_CUSTOM_KEYS = ( + SUBGRAPH_REGION, + SUBGRAPH_REGION_ROLE, + SUBGRAPH_REGION_PRESERVE_ORDER, +) + + +def subgraph( + name: str | None, + role: str | None = None, + *, + preserve_order: bool = False, +): + """Mark operations for outlining as explicit ``invoke_subgraph`` regions. + + Use this context inside code traced by ``minimal_fx_tracer``:: + + with subgraph(f"loss_chunk_{chunk_idx}", role="loss_chunk"): + loss = compute_chunk_loss(x) + + ``role`` gives regions with distinct ``name`` values a stable category for + downstream passes. For example, every iteration of a chunked loss can have + a unique name and share the ``"loss_chunk"`` role, allowing a later pass to + select only those outlined subgraphs for memory-policy tagging or + rematerialization. + + GraphTrainer copies the annotation from differentiable forward operations + to their generated backward operations. + ``apply_subgraph_region_annotations_pass`` outlines each contiguous + annotated segment. Forward and backward segments become separate subgraphs + with the same ``name`` and optional ``role``. Each outlined subgraph is an + Inductor fusion and buffer-reuse boundary. + """ + if name is None: + return nullcontext() + if not isinstance(name, str): + raise AssertionError( + f"expected subgraph region name to be str, got {type(name)}" + ) + custom: dict[str, Any] = {SUBGRAPH_REGION: name} + if role is not None: + if not isinstance(role, str): + raise AssertionError( + f"expected subgraph region role to be str, got {type(role)}" + ) + custom[SUBGRAPH_REGION_ROLE] = role + if preserve_order: + custom[SUBGRAPH_REGION_PRESERVE_ORDER] = True + return annotate(custom) + + +def _getattr_or_none(module: torch.fx.GraphModule, target: str) -> Any: + value: Any = module + missing = object() + for atom in target.split("."): + value = getattr(value, atom, missing) + if value is missing: + return None + return value + + +def _has_graph_module_arg(node: Node) -> bool: + gm = node.graph.owning_module + if gm is None: + return False + return any( + inp.op == "get_attr" + and isinstance(inp.target, str) + and isinstance(_getattr_or_none(gm, inp.target), torch.fx.GraphModule) + for inp in node.all_input_nodes + ) + + +def subgraph_region_key(node: Node) -> tuple[str, str, str, bool] | None: + if node.op in ("placeholder", "output", "get_attr"): + return None + + custom = node.meta.get("custom") + if not isinstance(custom, dict): + return None + region = custom.get(SUBGRAPH_REGION) + if region is None: + return None + if not isinstance(region, str): + raise AssertionError( + f"expected custom {SUBGRAPH_REGION} to be a str, got {type(region)}" + ) + role = custom.get(SUBGRAPH_REGION_ROLE) + if role is None: + role = "bw" if node.meta.get("autograd_backward") is True else "fw" + elif not isinstance(role, str): + raise AssertionError( + f"expected custom {SUBGRAPH_REGION_ROLE} to be a str, got {type(role)}" + ) + preserve_order = custom.get(SUBGRAPH_REGION_PRESERVE_ORDER, False) + if not isinstance(preserve_order, bool): + raise TypeError( + f"expected custom {SUBGRAPH_REGION_PRESERVE_ORDER} to be a bool, " + f"got {type(preserve_order)}" + ) + if ( + node.op == "call_function" + and isinstance(node.target, torch._ops.HigherOrderOperator) + and ( + node.target is torch.ops.higher_order.invoke_subgraph or not preserve_order + ) + ): + return None + if _has_graph_module_arg(node) and not preserve_order: + return None + return f"{region}_{role}", region, role, preserve_order + + +def collect_subgraph_region_groups( + graph: torch.fx.Graph, +) -> list[tuple[str, str, str, bool, list[Node]]]: + groups: list[tuple[str, str, str, bool, list[Node]]] = [] + current_key: tuple[str, str, str, bool] | None = None + current_nodes: list[Node] = [] + + def flush() -> None: + nonlocal current_key, current_nodes + if current_key is not None and len(current_nodes) > 1: + groups.append((*current_key, current_nodes)) + current_key = None + current_nodes = [] + + for node in list(graph.nodes): + if node.op == "get_attr" and current_key is not None and current_key[-1]: + continue + key = subgraph_region_key(node) + if key is None: + flush() + continue + if key != current_key: + flush() + current_key = key + current_nodes.append(node) + flush() + return groups + + +def _copy_placeholder_meta( + placeholder: fx.Node, input_node: fx.Node, owning_module: fx.GraphModule +) -> None: + if "val" in input_node.meta: + placeholder.meta.update(input_node.meta) + elif input_node.op == "get_attr" and isinstance(input_node.target, str): + placeholder.meta["val"] = attrgetter(input_node.target)(owning_module) + + +def _strip_custom_keys_from_meta(meta: dict[str, Any], keys: tuple[str, ...]) -> None: + if not keys: + return + custom = meta.get("custom") + if not isinstance(custom, dict): + return + + custom = custom.copy() + for key in keys: + custom.pop(key, None) + compile_with_inductor = custom.get("compile_with_inductor") + if isinstance(compile_with_inductor, dict): + compile_with_inductor = compile_with_inductor.copy() + for key in keys: + compile_with_inductor.pop(key, None) + custom["compile_with_inductor"] = compile_with_inductor + + if custom: + meta["custom"] = custom + else: + meta.pop("custom", None) + + +def _strip_subgraph_arg_annotations( + module: torch.fx.GraphModule, + nodes: list[Node], + keys: tuple[str, ...], +) -> None: + for node in nodes: + for input_node in node.all_input_nodes: + if input_node.op != "get_attr" or not isinstance(input_node.target, str): + continue + submodule = _getattr_or_none(module, input_node.target) + if not isinstance(submodule, torch.fx.GraphModule): + continue + for nested_module in submodule.modules(): + if not isinstance(nested_module, torch.fx.GraphModule): + continue + for subnode in nested_module.graph.nodes: + _strip_custom_keys_from_meta(subnode.meta, keys) + + +def _target_key(target): + if hasattr(target, "name") and hasattr(target, "overloadpacket"): + return str(target) + return repr(target) + + +def _meta_value_key(value): + if isinstance(value, torch.Tensor): + return ( + "tensor", + str(value.dtype), + tuple(_meta_value_key(dim) for dim in value.shape), + tuple(_meta_value_key(stride) for stride in value.stride()), + str(value.device), + value.requires_grad, + ) + if isinstance(value, (torch.SymInt, torch.SymFloat, torch.SymBool)): + return type(value).__name__, str(value) + if isinstance(value, (tuple, list)): + return type(value).__name__, tuple(_meta_value_key(item) for item in value) + if isinstance(value, dict): + return "dict", tuple( + sorted((repr(key), _meta_value_key(item)) for key, item in value.items()) + ) + if isinstance(value, (str, int, float, bool, type(None))): + return value + return repr(value) + + +def _arg_key(value, node_indices): + if isinstance(value, fx.Node): + return "node", node_indices[value] + if isinstance(value, (tuple, list)): + return type(value).__name__, tuple( + _arg_key(item, node_indices) for item in value + ) + if isinstance(value, dict): + return "dict", tuple( + sorted( + (repr(key), _arg_key(item, node_indices)) for key, item in value.items() + ) + ) + if isinstance(value, slice): + return ( + "slice", + _meta_value_key(value.start), + _meta_value_key(value.stop), + _meta_value_key(value.step), + ) + return _meta_value_key(value) + + +def _custom_key(custom): + if not isinstance(custom, dict): + return None + return tuple( + sorted( + (key, _meta_value_key(value)) + for key, value in custom.items() + if key != SUBGRAPH_REGION + ) + ) + + +def _node_meta_key(node): + return ( + _meta_value_key(node.meta.get("val")) if "val" in node.meta else None, + ( + _meta_value_key(node.meta.get("recompute")) + if "recompute" in node.meta + else None + ), + node.meta.get("autograd_backward", False), + _custom_key(node.meta.get("custom")), + ) + + +def _subgraph_structural_key(gm): + node_indices = {node: idx for idx, node in enumerate(gm.graph.nodes)} + return tuple( + ( + node.op, + _target_key(node.target), + _arg_key(node.args, node_indices), + _arg_key(node.kwargs, node_indices), + _node_meta_key(node), + ) + for node in gm.graph.nodes + ) + + +def _reuse_subgraph_module(module, region_node, canonical_target) -> None: + attr_node = region_node.args[0] + if not ( + isinstance(attr_node, Node) + and attr_node.op == "get_attr" + and isinstance(attr_node.target, str) + ): + return + + duplicate_target = attr_node.target + if duplicate_target == canonical_target: + return + + attr_node.target = canonical_target + region_node.args = (attr_node, canonical_target, *region_node.args[2:]) + if hasattr(module, duplicate_target): + delattr(module, duplicate_target) + + +def mark_invoke_subgraph( + graph: fx.Graph, + nodes: list[fx.Node], + *, + region_name_prefix: str, +) -> fx.Node: + """Outline FX nodes into an invoke_subgraph HOP and return the HOP node.""" + owning_module = graph.owning_module + if owning_module is None: + raise AssertionError("expected graph to have an owning_module") + if not nodes: + raise AssertionError("expected non-empty nodes") + + node_set = OrderedSet(nodes) + ordered_nodes = [node for node in graph.nodes if node in node_set] + if len(ordered_nodes) != len(node_set): + raise AssertionError("expected all nodes to belong to graph") + if any(node.op in ("placeholder", "output") for node in ordered_nodes): + raise AssertionError( + "expected invoke_subgraph nodes to exclude graph boundaries" + ) + + region_outputs = [ + node + for node in ordered_nodes + if any(user not in node_set for user in node.users) + ] + + subgraph = fx.Graph(owning_module) + env: dict[fx.Node, fx.Node] = {} + input_replacements: dict[fx.Node, Any] = {} + boundary_args: list[tuple[fx.Node, tuple[int, ...], Any]] = [] + + external_inputs: OrderedSet[fx.Node] = OrderedSet() + preserved_getattrs: OrderedSet[fx.Node] = OrderedSet() + + def collect_external_input(node: fx.Node) -> fx.Node: + if node not in node_set: + if ( + node.op == "get_attr" + and isinstance(node.target, str) + and isinstance(attrgetter(node.target)(owning_module), fx.GraphModule) + ): + preserved_getattrs.add(node) + else: + external_inputs.add(node) + return node + + for node in ordered_nodes: + fx.map_arg((node.args, node.kwargs), collect_external_input) + + node_order = {node: idx for idx, node in enumerate(graph.nodes)} + latest_input = max( + external_inputs, + key=lambda node: node_order[node], + default=None, + ) + first_external_user = min( + ( + user + for output_node in region_outputs + for user in output_node.users + if user not in node_set + ), + key=lambda node: node_order[node], + default=None, + ) + if ( + first_external_user is not None + and latest_input is not None + and node_order[latest_input] >= node_order[first_external_user] + ): + raise AssertionError("expected invoke_subgraph boundary to be acyclic") + + def add_boundary_arg( + input_node: fx.Node, path: tuple[int, ...], meta_val: Any + ) -> fx.Node: + placeholder = subgraph.placeholder(f"arg_{len(boundary_args)}") + _copy_placeholder_meta(placeholder, input_node, owning_module) + if path: + placeholder.meta["val"] = meta_val + boundary_args.append((input_node, path, meta_val)) + return placeholder + + def make_input_replacement( + input_node: fx.Node, value: Any, path: tuple[int, ...] = () + ) -> Any: + if isinstance(value, (tuple, list)): + return type(value)( + make_input_replacement(input_node, item, (*path, idx)) + for idx, item in enumerate(value) + ) + return add_boundary_arg(input_node, path, value) + + for input_node in external_inputs: + value = input_node.meta.get("val") + if isinstance(value, (tuple, list)): + input_replacements[input_node] = make_input_replacement(input_node, value) + else: + input_replacements[input_node] = add_boundary_arg(input_node, (), value) + + def load_arg(node: fx.Node) -> Any: + if node in env: + return env[node] + if node in node_set: + raise AssertionError("expected invoke_subgraph nodes to be topological") + if node in preserved_getattrs: + if not isinstance(node.target, str): + raise AssertionError("expected get_attr target to be a string") + get_attr_node = subgraph.get_attr(node.target) + get_attr_node.meta.update(node.meta) + env[node] = get_attr_node + return get_attr_node + return input_replacements[node] + + for node in ordered_nodes: + env[node] = subgraph.node_copy(node, load_arg) + + subgraph_outputs = tuple(env[node] for node in region_outputs) + out = subgraph.output(subgraph_outputs) + out.meta["val"] = tuple(node.meta.get("val") for node in region_outputs) + subgraph.lint() + + subgraph_module = _LazyGraphModule(owning_module, subgraph) + first_name = ordered_nodes[0].name + last_name = ordered_nodes[-1].name + region_name = f"{region_name_prefix}_{first_name}_{last_name}" + subgraph_attr_name = f"{region_name}_0" + setattr(owning_module, subgraph_attr_name, subgraph_module) + + if latest_input is None or node_order[latest_input] < node_order[ordered_nodes[0]]: + with graph.inserting_before(ordered_nodes[0]): + get_subgraph = graph.get_attr(subgraph_attr_name) + else: + with graph.inserting_after(latest_input): + get_subgraph = graph.get_attr(subgraph_attr_name) + + outer_args: list[fx.Node] = [] + insert_after = get_subgraph + + def make_outer_arg( + input_node: fx.Node, path: tuple[int, ...], meta_val: Any + ) -> fx.Node: + nonlocal insert_after + source = input_node + for idx in path: + with graph.inserting_after(insert_after): + source = graph.call_function( + getitem, + args=(source, idx), + name=f"{input_node.name}_{region_name_prefix}_arg_{len(outer_args)}", + ) + insert_after = source + if path: + source.meta["val"] = meta_val + return source + + for input_node, path, meta_val in boundary_args: + outer_args.append(make_outer_arg(input_node, path, meta_val)) + + with graph.inserting_after(insert_after): + region_node = graph.call_function( + torch.ops.higher_order.invoke_subgraph, + args=(get_subgraph, subgraph_attr_name, *outer_args), + name=region_name, + ) + + replacements: list[fx.Node] = [] + if len(region_outputs) == 0: + region_node.meta["val"] = () + else: + region_node.meta["val"] = tuple(node.meta.get("val") for node in region_outputs) + insert_after = region_node + for idx, output_node in enumerate(region_outputs): + with graph.inserting_after(insert_after): + replacement = graph.call_function( + getitem, + args=(region_node, idx), + name=f"{output_node.name}_{region_name_prefix}", + ) + replacement.meta = output_node.meta.copy() + replacement.meta.pop("eager_input_vals", None) + _strip_custom_keys_from_meta(replacement.meta, _SUBGRAPH_REGION_CUSTOM_KEYS) + replacements.append(replacement) + insert_after = replacement + + for output_node, replacement in zip(region_outputs, replacements, strict=True): + for user in list(output_node.users): + if user not in node_set: + user.replace_input_with(output_node, replacement) + + for node in reversed(ordered_nodes): + graph.erase_node(node) + graph.lint() + + return region_node + + +def _record_subgraph_region( + module: torch.fx.GraphModule, + region_node: Node, + region: str, + region_id: str, + region_role: str, + preserve_order: bool, +) -> None: + region_node.meta[SUBGRAPH_REGION] = region + custom = region_node.meta.setdefault("custom", {}) + custom["subgraph_region_id"] = region_id + custom["subgraph_region_role"] = region_role + nested_config = None + if preserve_order: + from torch._higher_order_ops.invoke_subgraph import NestedCompileRegionOptions + + nested_config = NestedCompileRegionOptions( + inductor_config_patches={ + "reorder_for_locality": False, + "reorder_for_peak_memory": False, + "reorder_for_compute_comm_overlap": False, + # Without a later scheduling pass, pre-fusion lifetime metadata + # would be stale after fusion. + "fusion_memory_timeline_peak_allowed_increase_mb": None, + "aten_distributed_optimizations.enable_simple_overlap": False, + "aten_distributed_optimizations.enable_overlap_scheduling": False, + } + ) + custom["nested_region_config"] = nested_config + get_subgraph = region_node.args[0] + if not ( + isinstance(get_subgraph, Node) + and get_subgraph.op == "get_attr" + and isinstance(get_subgraph.target, str) + ): + return + submod = getattr(module, get_subgraph.target, None) + if isinstance(submod, torch.fx.GraphModule): + submod.meta[SUBGRAPH_REGION] = region + submod_custom = submod.meta.setdefault("custom", {}) + if not isinstance(submod_custom, dict): + submod_custom = {} + submod.meta["custom"] = submod_custom + submod_custom["subgraph_region_id"] = region_id + submod_custom["subgraph_region_role"] = region_role + if nested_config is not None: + submod.meta["nested_region_config"] = nested_config + + +def apply_subgraph_region_annotations_pass( + gm: torch.fx.GraphModule, + example_inputs: tuple | None = None, +) -> torch.fx.GraphModule: + del example_inputs + + outlined_regions = 0 + for module in list(gm.modules()): + if not isinstance(module, torch.fx.GraphModule): + continue + outlined_subgraphs = {} + groups = collect_subgraph_region_groups(module.graph) + if not groups: + continue + for region, region_id, region_role, preserve_order, nodes in groups: + if preserve_order: + _strip_subgraph_arg_annotations( + module, nodes, _SUBGRAPH_REGION_CUSTOM_KEYS + ) + region_node = mark_invoke_subgraph( + module.graph, + nodes, + region_name_prefix=f"subgraph_region_{outlined_regions}", + ) + _record_subgraph_region( + module, + region_node, + region, + region_id, + region_role, + preserve_order, + ) + attr_node = region_node.args[0] + if ( + isinstance(attr_node, Node) + and attr_node.op == "get_attr" + and isinstance(attr_node.target, str) + ): + submod = getattr(module, attr_node.target, None) + if isinstance(submod, torch.fx.GraphModule): + subgraph_key = _subgraph_structural_key(submod) + canonical_target = outlined_subgraphs.setdefault( + subgraph_key, attr_node.target + ) + _reuse_subgraph_module(module, region_node, canonical_target) + outlined_regions += 1 + module.graph.lint() + module.recompile() + + return gm diff --git a/torchtitan/experiments/graph_trainer/tests/test_passes.py b/torchtitan/experiments/graph_trainer/tests/test_passes.py index 2507170406..c4fa2bf327 100644 --- a/torchtitan/experiments/graph_trainer/tests/test_passes.py +++ b/torchtitan/experiments/graph_trainer/tests/test_passes.py @@ -11,6 +11,7 @@ from unittest.mock import patch import torch +from torch._decomp import get_decompositions from torch._functorch.aot_autograd import aot_compile_joint_with_descriptors from torch._guards import tracing from torch._inductor.fx_passes.bucketing import ( @@ -19,6 +20,7 @@ from torch.cuda._graph_annotations import _is_tools_id_unavailable from torch.fx.experimental.proxy_tensor import make_fx from torch.fx.experimental.symbolic_shapes import ShapeEnv +from torch.fx.passes.fake_tensor_prop import FakeTensorProp from torch.fx.traceback import preserve_node_meta from torch.testing._internal.common_fsdp import FSDPTest from torch.testing._internal.common_utils import TestCase @@ -44,6 +46,9 @@ is_cudagraphable, is_full_cudagraphable, ) +from torchtitan.experiments.graph_trainer.decompositions import ( + apply_decompositions_pass, +) from torchtitan.experiments.graph_trainer.ep_chunk_pass import ( _chunk_copied_meta, _materialize_symint_arg, @@ -88,9 +93,11 @@ run_traced, ) from torchtitan.experiments.graph_trainer.memory_policy import ( + _backward_side_nodes, _default_memory_policy_pass, _make_default_memory_policy, _make_full_memory_policy, + tag_min_cut_saved_values, tag_sac_policy, tag_with_memory_policy_pass, validate_memory_policy_config, @@ -109,6 +116,11 @@ remove_identity_view_pass, ) from torchtitan.experiments.graph_trainer.simple_fsdp import data_parallel +from torchtitan.experiments.graph_trainer.subgraph_regions import ( + apply_subgraph_region_annotations_pass, + SUBGRAPH_REGION, + SUBGRAPH_REGION_ROLE, +) from torchtitan.experiments.graph_trainer.tests.test_cpu_offload import ( # noqa: F401 TestCpuOffloadPass, ) @@ -1829,6 +1841,207 @@ def test_same_layer_nodes_all_recomputed(self): ) +class TestMinCutMemoryPolicy(TestCase): + @staticmethod + def _config(): + return SimpleNamespace( + compile=GraphTrainerCompileConfig(memory_policy="min_cut") + ) + + @staticmethod + def _fake_prop(gm, *inputs): + with torch._subclasses.FakeTensorMode() as fake_mode: + fake_inputs = [ + torch.empty(shape, device="cuda", dtype=dtype) + for shape, dtype in inputs + ] + FakeTensorProp(gm, mode=fake_mode).propagate_dont_convert_inputs( + *fake_inputs + ) + + @staticmethod + def _recomputed_nodes(gm): + return [node for node in gm.graph.nodes if node.name.endswith("_recomputed")] + + @staticmethod + def _log_softmax_decomposition_table(): + return get_decompositions([torch.ops.aten._log_softmax.default]) + + def test_view_cut_saves_its_base(self): + graph = torch.fx.Graph() + x = graph.placeholder("x") + weight = graph.placeholder("weight") + grad = graph.placeholder("grad") + mm = graph.call_function(torch.ops.aten.mm.default, args=(x, weight)) + view = graph.call_function(torch.ops.aten.view.default, args=(mm, [4, 4])) + bwd = graph.call_function(torch.ops.aten.mm.default, args=(view, grad)) + bwd.meta["autograd_backward"] = True + graph.output(bwd) + gm = torch.fx.GraphModule(torch.nn.Module(), graph) + + tag_min_cut_saved_values(gm, _backward_side_nodes(gm), {view}) + + self.assertEqual(mm.meta["recompute"], CheckpointPolicy.MUST_SAVE) + self.assertEqual(view.meta["recompute"], CheckpointPolicy.MUST_RECOMPUTE) + + def test_applies_to_whole_graph(self): + graph = torch.fx.Graph() + x = graph.placeholder("x") + a = graph.call_function(torch.ops.aten.sin.default, args=(x,)) + b = graph.call_function(torch.ops.aten.cos.default, args=(a,)) + loss = graph.call_function(torch.ops.aten.sum.default, args=(b,)) + bwd = graph.call_function(torch.ops.aten.neg.default, args=(b,)) + bwd.meta["autograd_backward"] = True + graph.output((loss, bwd)) + gm = torch.fx.GraphModule(torch.nn.Module(), graph) + self._fake_prop(gm, ((64, 64), torch.float32)) + + tag_with_memory_policy_pass(gm, config=self._config()) + self.assertTrue( + any( + node.meta.get("recompute") == CheckpointPolicy.MUST_RECOMPUTE + for node in gm.graph.nodes + ) + ) + self.assertEqual(len(self._recomputed_nodes(gm)), 0) + selective_activation_remat_pass(gm) + + self.assertGreaterEqual(len(self._recomputed_nodes(gm)), 1) + + def test_decomposition_is_a_standalone_pass_before_min_cut(self): + graph = torch.fx.Graph() + x = graph.placeholder("x") + grad = graph.placeholder("grad") + log_probs = graph.call_function( + torch.ops.aten._log_softmax.default, args=(x, -1, False) + ) + loss = graph.call_function(torch.ops.aten.sum.default, args=(log_probs,)) + bwd = graph.call_function( + torch.ops.aten._log_softmax_backward_data.default, + args=(grad, log_probs, -1, torch.float32), + ) + bwd.meta["autograd_backward"] = True + graph.output((loss, bwd)) + gm = torch.fx.GraphModule(torch.nn.Module(), graph) + self._fake_prop(gm, ((128, 1024), torch.float32), ((128, 1024), torch.float32)) + + apply_decompositions_pass( + gm, + decomposition_table=self._log_softmax_decomposition_table(), + ) + self.assertFalse( + any( + node.target == torch.ops.aten._log_softmax.default + for node in gm.graph.nodes + ) + ) + self.assertEqual(len(self._recomputed_nodes(gm)), 0) + + tag_with_memory_policy_pass(gm, config=self._config()) + bwd = next( + node + for node in gm.graph.nodes + if node.target == torch.ops.aten._log_softmax_backward_data.default + ) + self.assertIsInstance(bwd.args[1], torch.fx.Node) + self.assertEqual(bwd.args[1].meta["recompute"], CheckpointPolicy.MUST_SAVE) + selective_activation_remat_pass(gm) + + bwd = next( + node + for node in gm.graph.nodes + if node.target == torch.ops.aten._log_softmax_backward_data.default + ) + self.assertIsInstance(bwd.args[1], torch.fx.Node) + self.assertFalse(bwd.args[1].name.endswith("_recomputed")) + + def test_min_cut_policy_respects_existing_checkpoint_policy(self): + graph = torch.fx.Graph() + x = graph.placeholder("x") + saved = graph.call_function(torch.ops.aten.sin.default, args=(x,)) + recompute = graph.call_function(torch.ops.aten.cos.default, args=(saved,)) + loss = graph.call_function(torch.ops.aten.sum.default, args=(recompute,)) + bwd = graph.call_function(torch.ops.aten.neg.default, args=(recompute,)) + bwd.meta["autograd_backward"] = True + graph.output((loss, bwd)) + gm = torch.fx.GraphModule(torch.nn.Module(), graph) + self._fake_prop(gm, ((64, 64), torch.float32)) + saved.meta["recompute"] = CheckpointPolicy.MUST_SAVE + recompute.meta["recompute"] = CheckpointPolicy.PREFER_RECOMPUTE + + tag_with_memory_policy_pass(gm, config=self._config()) + + self.assertEqual(saved.meta["recompute"], CheckpointPolicy.MUST_SAVE) + self.assertIn( + recompute.meta["recompute"], + (CheckpointPolicy.PREFER_RECOMPUTE, CheckpointPolicy.MUST_RECOMPUTE), + ) + selective_activation_remat_pass(gm) + recomputed_targets = {node.target for node in self._recomputed_nodes(gm)} + self.assertIn(torch.ops.aten.cos.default, recomputed_targets) + self.assertNotIn(torch.ops.aten.sin.default, recomputed_targets) + + def test_explicit_subgraph_decomposition_and_min_cut_policy(self): + graph = torch.fx.Graph() + x = graph.placeholder("x") + grad = graph.placeholder("grad") + log_probs = graph.call_function( + torch.ops.aten._log_softmax.default, args=(x, -1, False) + ) + loss = graph.call_function(torch.ops.aten.sum.default, args=(log_probs,)) + bwd = graph.call_function( + torch.ops.aten._log_softmax_backward_data.default, + args=(grad, log_probs, -1, torch.float32), + ) + graph.output((loss, bwd)) + gm = torch.fx.GraphModule(torch.nn.Module(), graph) + self._fake_prop(gm, ((128, 1024), torch.float32), ((128, 1024), torch.float32)) + + for node in (log_probs, loss, bwd): + node.meta.setdefault("custom", {}) + node.meta["custom"][SUBGRAPH_REGION] = "region" + node.meta["custom"][SUBGRAPH_REGION_ROLE] = "fw_bw_grad_accum" + bwd.meta["autograd_backward"] = True + + apply_subgraph_region_annotations_pass(gm) + apply_decompositions_pass( + gm, + decomposition_table=self._log_softmax_decomposition_table(), + recurse=True, + apply_to_root=False, + ) + submods = [ + module + for module in gm.modules() + if isinstance(module, torch.fx.GraphModule) and module is not gm + ] + self.assertEqual(len(submods), 1) + submod = submods[0] + tag_with_memory_policy_pass(submod, config=self._config()) + self.assertFalse( + any( + node.target == torch.ops.aten._log_softmax.default + for node in submod.graph.nodes + ) + ) + bwd = next( + node + for node in submod.graph.nodes + if node.target == torch.ops.aten._log_softmax_backward_data.default + ) + self.assertIsInstance(bwd.args[1], torch.fx.Node) + self.assertEqual(bwd.args[1].meta["recompute"], CheckpointPolicy.MUST_SAVE) + + selective_activation_remat_pass(submod) + bwd = next( + node + for node in submod.graph.nodes + if node.target == torch.ops.aten._log_softmax_backward_data.default + ) + self.assertIsInstance(bwd.args[1], torch.fx.Node) + self.assertFalse(bwd.args[1].name.endswith("_recomputed")) + + class TestBucketingPrefetchOrder(FSDPTest): """Guard that SAC + bucketing produces correct all_gather prefetch order. diff --git a/torchtitan/experiments/graph_trainer/tests/test_profiler.py b/torchtitan/experiments/graph_trainer/tests/test_profiler.py index 74740934de..dee8e1f4ec 100644 --- a/torchtitan/experiments/graph_trainer/tests/test_profiler.py +++ b/torchtitan/experiments/graph_trainer/tests/test_profiler.py @@ -4,22 +4,18 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -import gzip import json import os import tempfile import unittest +from typing import Any from unittest.mock import patch import torch -from torch.cuda._annotate_cuda_graph_trace import ( # pyrefly: ignore[missing-import] - annotate_trace, -) from torch.cuda._graph_annotations import _is_tools_id_unavailable from torch.testing._internal.common_utils import run_tests, TestCase from torchtitan.distributed.cudagraph import ( - cudagraph_annotate_trace_post_processor, cudagraph_teardown, get_cudagraph_annotations, ) @@ -35,7 +31,7 @@ apply_graph_passes, construct_default_graph_passes, ) -from torchtitan.tools.profiler import Profiler +from torchtitan.tools.profiler import _EXPORT_SUPPORTS_ANNOTATIONS, Profiler @unittest.skipUnless(torch.cuda.is_available(), "CUDA not available") @@ -48,6 +44,8 @@ def test_profiler_trace_has_module_fqn_annotations(self): ``module_fqn`` fields on graphed kernel events.""" if _is_tools_id_unavailable(): self.skipTest("cudaGraphNodeGetToolsId not available") + if not _EXPORT_SUPPORTS_ANNOTATIONS: + self.skipTest("export_chrome_trace has no cuda_graph_annotations argument") # Simple model with annotated submodules. class FFN(torch.nn.Module): @@ -124,14 +122,11 @@ def fwd_bwd_step(inputs, labels): with tempfile.NamedTemporaryFile(suffix=".json", delete=False) as f: trace_path = f.name - prof.export_chrome_trace(trace_path) + prof.export_chrome_trace(trace_path, cuda_graph_annotations=annotations) with open(trace_path) as f: trace = json.load(f) - count = annotate_trace(trace, annotations) - self.assertGreater(count, 0, "annotate_trace matched 0 events") - # Verify module_fqn fields appear on graphed kernel events. # Since minimal_fx_tracer traces fwd+bwd into a single graph, # backward kernels (e.g. layer_norm_backward) should also carry @@ -176,23 +171,34 @@ def fwd_bwd_step(inputs, labels): cudagraph_teardown() -class TestTracePostProcessor(TestCase): - """Verify CUDA graph annotations are applied to every profiler trace.""" +class TestTraceAnnotationExport(TestCase): + """Verify CUDA graph annotations reach every profiler trace the Profiler writes.""" - def test_post_processor_called_with_trace_path(self): - """The CUDA graph post-processor receives the exported trace path.""" + ANNOTATIONS = {42: [{_MODULE_FQN: "layers.0.attention.wq"}]} - calls: list[tuple[str, bool]] = [] + def _run_profiler(self, supports_annotations: bool) -> list[tuple[str, Any]]: + """Drive one profile cycle, returning (path, cuda_graph_annotations) per export.""" + calls: list[tuple[str, Any]] = [] - def record_call(trace_path: str) -> None: - calls.append((trace_path, os.path.exists(trace_path))) + def record_export(self_prof, path, *args, **kwargs): + calls.append((path, kwargs.get("cuda_graph_annotations"))) with ( tempfile.TemporaryDirectory() as tmp, patch("torch.distributed.get_rank", return_value=0), patch( - "torchtitan.tools.profiler.cudagraph_annotate_trace_post_processor", - side_effect=record_call, + "torchtitan.tools.profiler.get_cudagraph_annotations", + return_value=self.ANNOTATIONS, + ), + patch( + "torchtitan.tools.profiler._EXPORT_SUPPORTS_ANNOTATIONS", + supports_annotations, + ), + patch.object( + torch.profiler.profile, + "export_chrome_trace", + autospec=True, + side_effect=record_export, ), ): config = Profiler.Config( @@ -208,53 +214,24 @@ def record_call(trace_path: str) -> None: for _ in range(4): profiler.step() - self.assertEqual(len(calls), 1, f"Expected 1 call, got {calls}") - path, existed = calls[0] - # Profiler exports gzip-compressed traces (.json.gz) since #3483. + self.assertEqual(len(calls), 1, f"Expected 1 export, got {calls}") + return calls + + def test_annotations_baked_into_export(self): + """The trace handler hands the captured annotations to the export rather than + joining them onto the written file afterwards.""" + path, passed = self._run_profiler(supports_annotations=True)[0] + # Profiler exports gzip-compressed traces (.json.gz) since #3483; the exporter + # keys compression off that suffix and bakes the annotations in as it writes. self.assertTrue(path.endswith("rank0_trace.json.gz")) - self.assertTrue(existed, f"Trace file {path} did not exist when callback ran") - - def test_annotate_post_processor_round_trips_gzip_trace(self): - """The cudagraph trace post-processor must read and write the - gzip-compressed (.json.gz) traces the Profiler produces since #3483. - A plain ``open``/``json.load`` would raise on the gzip bytes.""" - - graph_node_id = 42 - trace = { - "traceEvents": [ - { - "name": "some_kernel", - "tid": 1, - "ts": 100, - "args": {"graph node id": graph_node_id}, - } - ] - } - - # Seed annotations so the post-processor does real work instead of - # returning early on an empty annotation map. - annotations = get_cudagraph_annotations() - saved_annotations = annotations.copy() - annotations.clear() - annotations[graph_node_id] = [{_MODULE_FQN: "layers.0.attention.wq"}] - try: - with tempfile.TemporaryDirectory() as tmp: - trace_path = os.path.join(tmp, "rank0_trace.json.gz") - with gzip.open(trace_path, "wt") as f: - json.dump(trace, f) - - # Must not raise on the gzip-compressed input. - cudagraph_annotate_trace_post_processor(trace_path) - - # And the written-back trace must remain valid gzip JSON. - with gzip.open(trace_path, "rt") as f: - annotated = json.load(f) - finally: - annotations.clear() - annotations.update(saved_annotations) - - fqns = {e.get("args", {}).get(_MODULE_FQN) for e in annotated["traceEvents"]} - self.assertIn("layers.0.attention.wq", fqns) + self.assertEqual(passed, self.ANNOTATIONS) + + def test_export_still_runs_without_annotation_support(self): + """On a torch whose export_chrome_trace predates cuda_graph_annotations the + trace is still written, just without them.""" + path, passed = self._run_profiler(supports_annotations=False)[0] + self.assertTrue(path.endswith("rank0_trace.json.gz")) + self.assertIsNone(passed) if __name__ == "__main__": diff --git a/torchtitan/experiments/graph_trainer/tests/test_subgraph_regions.py b/torchtitan/experiments/graph_trainer/tests/test_subgraph_regions.py new file mode 100644 index 0000000000..03bbb9afce --- /dev/null +++ b/torchtitan/experiments/graph_trainer/tests/test_subgraph_regions.py @@ -0,0 +1,189 @@ +# 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 torch +import torch.fx as fx +from torch.fx.experimental.proxy_tensor import make_fx +from torch.fx.experimental.symbolic_shapes import ShapeEnv +from torch.fx.traceback import preserve_node_meta +from torch.testing._internal.common_utils import run_tests, TestCase + +from torchtitan.experiments.graph_trainer.make_fx_tracer import minimal_fx_tracer +from torchtitan.experiments.graph_trainer.subgraph_regions import ( + apply_subgraph_region_annotations_pass, + subgraph, + SUBGRAPH_REGION, + SUBGRAPH_REGION_ROLE, +) + + +def _annotate_region(nodes, region): + for node in nodes: + node.meta.setdefault("custom", {}) + node.meta["custom"][SUBGRAPH_REGION] = region + node.meta["custom"][SUBGRAPH_REGION_ROLE] = "loss_chunk" + + +def _invoke_subgraph_nodes(gm): + return list( + gm.graph.find_nodes( + op="call_function", + target=torch.ops.higher_order.invoke_subgraph, + ) + ) + + +def _subgraph_modules(gm): + return [ + module + for module in gm.modules() + if isinstance(module, fx.GraphModule) and module is not gm + ] + + +class TestSubgraphRegions(TestCase): + def test_dedup_is_scoped_to_owning_module(self): + def make_submodule(): + graph = fx.Graph() + x = graph.placeholder("x") + sin = graph.call_function(torch.ops.aten.sin.default, args=(x,)) + cos = graph.call_function(torch.ops.aten.cos.default, args=(sin,)) + graph.output(cos) + gm = fx.GraphModule(torch.nn.Module(), graph) + _annotate_region((sin, cos), "chunk") + return gm + + root = torch.nn.Module() + root.left = make_submodule() + root.right = make_submodule() + graph = fx.Graph() + x = graph.placeholder("x") + left = graph.call_module("left", args=(x,)) + right = graph.call_module("right", args=(x,)) + graph.output((left, right)) + gm = fx.GraphModule(root, graph) + + apply_subgraph_region_annotations_pass(gm) + + for module in (gm.left, gm.right): + invoke_node = _invoke_subgraph_nodes(module)[0] + attr_node = invoke_node.args[0] + self.assertTrue(hasattr(module, attr_node.target)) + + def test_subgraph_context_outlines_forward_and_backward_regions(self): + def train_step(x): + x = x.detach().requires_grad_() + with subgraph("loss_chunk_0", role="loss_chunk"): + y = torch.ops.aten.sin.default(x) + loss = torch.ops.aten.sum.default(y) + (grad,) = torch.autograd.grad(loss, (x,)) + return loss.detach(), grad + + x = torch.randn(4) + gm = minimal_fx_tracer(train_step)(x).gm + + apply_subgraph_region_annotations_pass(gm) + + invoke_nodes = _invoke_subgraph_nodes(gm) + self.assertEqual(len(invoke_nodes), 2) + for node in invoke_nodes: + self.assertEqual(node.meta[SUBGRAPH_REGION], "loss_chunk_0_loss_chunk") + self.assertEqual(node.meta["custom"]["subgraph_region_id"], "loss_chunk_0") + self.assertEqual(node.meta["custom"]["subgraph_region_role"], "loss_chunk") + + subgraph_modules = _subgraph_modules(gm) + self.assertEqual(len(subgraph_modules), 2) + self.assertEqual( + { + any(node.meta.get("autograd_backward") for node in module.graph.nodes) + for module in subgraph_modules + }, + {False, True}, + ) + torch.testing.assert_close(gm(x), train_step(x)) + + def test_preserve_order_outlines_higher_order_operator(self): + def true_fn(x): + return x.sin() + + def false_fn(x): + return x.cos() + + def f(pred, x): + with subgraph("ordered", preserve_order=True): + x = x + 1 + x = torch.cond(pred, true_fn, false_fn, (x,)) + return x * 2 + + pred = torch.tensor(True) + x = torch.randn(4) + with preserve_node_meta(): + gm = make_fx(f)(pred, x) + + apply_subgraph_region_annotations_pass(gm) + + invoke_nodes = _invoke_subgraph_nodes(gm) + self.assertEqual(len(invoke_nodes), 1) + nested_config = invoke_nodes[0].meta["custom"]["nested_region_config"] + self.assertEqual( + nested_config.inductor_config_patches, + { + "reorder_for_locality": False, + "reorder_for_peak_memory": False, + "reorder_for_compute_comm_overlap": False, + "fusion_memory_timeline_peak_allowed_increase_mb": None, + "aten_distributed_optimizations.enable_simple_overlap": False, + "aten_distributed_optimizations.enable_overlap_scheduling": False, + }, + ) + self.assertEqual(gm(pred, x), f(pred, x)) + + def test_structurally_identical_subgraph_regions_reuse_submodule(self): + graph = fx.Graph() + x = graph.placeholder("x") + a0 = graph.call_function(torch.ops.aten.sin.default, args=(x,)) + b0 = graph.call_function(torch.ops.aten.cos.default, args=(a0,)) + a1 = graph.call_function(torch.ops.aten.sin.default, args=(x,)) + b1 = graph.call_function(torch.ops.aten.cos.default, args=(a1,)) + out = graph.call_function(torch.ops.aten.add.Tensor, args=(b0, b1)) + graph.output((out,)) + gm = fx.GraphModule(torch.nn.Module(), graph) + _annotate_region((a0, b0), "chunk_0") + _annotate_region((a1, b1), "chunk_1") + symbolic_value = torch.empty(ShapeEnv().create_unbacked_symint(), device="meta") + for node in (a0, b0, a1, b1): + node.meta["val"] = symbolic_value + + apply_subgraph_region_annotations_pass(gm) + + invoke_nodes = _invoke_subgraph_nodes(gm) + self.assertEqual(len(invoke_nodes), 2) + self.assertEqual(invoke_nodes[0].args[1], invoke_nodes[1].args[1]) + self.assertEqual(len(_subgraph_modules(gm)), 1) + + def test_structurally_different_subgraph_regions_keep_submodules(self): + graph = fx.Graph() + x = graph.placeholder("x") + a0 = graph.call_function(torch.ops.aten.sin.default, args=(x,)) + b0 = graph.call_function(torch.ops.aten.cos.default, args=(a0,)) + a1 = graph.call_function(torch.ops.aten.sin.default, args=(x,)) + b1 = graph.call_function(torch.ops.aten.neg.default, args=(a1,)) + out = graph.call_function(torch.ops.aten.add.Tensor, args=(b0, b1)) + graph.output((out,)) + gm = fx.GraphModule(torch.nn.Module(), graph) + _annotate_region((a0, b0), "chunk_0") + _annotate_region((a1, b1), "chunk_1") + + apply_subgraph_region_annotations_pass(gm) + + invoke_nodes = _invoke_subgraph_nodes(gm) + self.assertEqual(len(invoke_nodes), 2) + self.assertNotEqual(invoke_nodes[0].args[1], invoke_nodes[1].args[1]) + self.assertEqual(len(_subgraph_modules(gm)), 2) + + +if __name__ == "__main__": + run_tests() diff --git a/torchtitan/experiments/rl/README.md b/torchtitan/experiments/rl/README.md index 5df4088888..cad2b5f4d7 100644 --- a/torchtitan/experiments/rl/README.md +++ b/torchtitan/experiments/rl/README.md @@ -71,7 +71,9 @@ def my_experiment() -> Controller.Config: return Controller.Config( model_spec=..., rollouter=MyRollouter.Config(), - renderer=RendererConfig(...), + renderer=RenderersLibraryConfig( + renderers_config=Qwen3RendererConfig(enable_thinking=False) + ), trainer=PolicyTrainer.Config(...), generator=VLLMGenerator.Config(...), ) @@ -110,7 +112,7 @@ uv venv --python 3.12 titan-rl source titan-rl/bin/activate ``` -1. Install Monarch, TorchStore, and Renderers from main: +1. Install Monarch, TorchStore, and Renderers: ```bash uv pip install -r torchtitan/experiments/rl/requirements.txt uv pip install --no-deps "git+https://github.com/meta-pytorch/torchstore.git@main" diff --git a/torchtitan/experiments/rl/actors/rollout_worker.py b/torchtitan/experiments/rl/actors/rollout_worker.py index a292c3dbab..6084769204 100644 --- a/torchtitan/experiments/rl/actors/rollout_worker.py +++ b/torchtitan/experiments/rl/actors/rollout_worker.py @@ -11,8 +11,9 @@ from typing import Any from monarch.actor import Actor, concurrent_endpoint -from torchtitan.experiments.rl.renderer import RendererConfig +from torchtitan.components.tokenizer import HuggingFaceTokenizer +from torchtitan.experiments.rl.renderer import RendererConfig from torchtitan.experiments.rl.rollout.rollouter import RolloutWorker from torchtitan.experiments.rl.rollout.types import RolloutGroup from torchtitan.observability import structured_logger as sl @@ -36,10 +37,12 @@ def __init__( async def setup_async( self, *, + tokenizer_config: HuggingFaceTokenizer.Config, renderer_config: RendererConfig, hf_assets_path: str, ) -> None: await self._worker.setup_async( + tokenizer_config=tokenizer_config, renderer_config=renderer_config, hf_assets_path=hf_assets_path, ) diff --git a/torchtitan/experiments/rl/controller.py b/torchtitan/experiments/rl/controller.py index 2ece881348..ab7dad0022 100644 --- a/torchtitan/experiments/rl/controller.py +++ b/torchtitan/experiments/rl/controller.py @@ -104,6 +104,7 @@ from monarch.actor import ProcMesh, this_host from monarch.spmd import setup_torch_elastic_env_async +from torchtitan.components.tokenizer import HuggingFaceTokenizer from torchtitan.config import CompileConfig, Configurable from torchtitan.experiments.rl.actors.generator import SamplingConfig, VLLMGenerator from torchtitan.experiments.rl.actors.trainer import PolicyTrainer @@ -297,8 +298,14 @@ class Config(Configurable.Config): """The rollouter: its datasets, envs, and rubric.""" # TODO: support multiple rollouters for data mixing. + tokenizer: HuggingFaceTokenizer.Config = field( + default_factory=HuggingFaceTokenizer.Config + ) + """Tokenizer loaded from `hf_assets_path`.""" + renderer: RendererConfig - """Message-to-token renderer config.""" + """The model's chat template; renders messages to token ids and parses completions + back. E.g. `RenderersLibraryConfig(renderers_config=Qwen3RendererConfig(enable_thinking=False))`.""" rollout_recorder: RolloutSampleRecorder.Config = field( default_factory=RolloutSampleRecorder.Config @@ -424,7 +431,8 @@ def __init__(self, config: Config): log_dir=config.dump_folder, job_config=config.to_dict(), ) - self.renderer = config.renderer.build(tokenizer_path=config.hf_assets_path) + self.tokenizer = config.tokenizer.build(tokenizer_path=config.hf_assets_path) + self.renderer = config.renderer.build(tokenizer=self.tokenizer) # Carry the base seed and renderer stop tokens on the sampling config so # the generator reads them off each request; the rollouter offsets the @@ -434,10 +442,6 @@ def __init__(self, config: Config): seed=config.generator.debug.seed, stop_token_ids=list(self.renderer.get_stop_token_ids()), ) - # TODO: pass our own tokenizer to the renderer and read pad/eos off it - # once `renderers` supports bring-your-own-tokenizer - # (https://github.com/PrimeIntellect-ai/renderers/pull/70). - # Until then, reach into the renderer's tokenizer for the pad id (eos doubles as pad). self._rollouter: Rollouter = config.rollouter.build() self.rollout_recorder = config.rollout_recorder.build( dump_dir=config.dump_folder @@ -636,6 +640,7 @@ async def setup_async( ) await self._rollouter.setup_async( + tokenizer_config=config.tokenizer, renderer_config=config.renderer, hf_assets_path=config.hf_assets_path, ) @@ -809,7 +814,7 @@ async def run(self) -> None: max_context_length=self.config.trainer.training.max_context_length, num_prompts_per_train_step=async_loop.num_prompts_per_train_step, dp_degree=self.trainer_dp_degree, - pad_id=self.renderer._tokenizer.eos_token_id, + pad_id=self.tokenizer.eos_id, ) # training_batch_queue diff --git a/torchtitan/experiments/rl/environment/token.py b/torchtitan/experiments/rl/environment/token.py index 3a528a4b88..29d1dba86c 100644 --- a/torchtitan/experiments/rl/environment/token.py +++ b/torchtitan/experiments/rl/environment/token.py @@ -161,6 +161,7 @@ async def step(self, completion: Completion) -> TokenEnvOutput: parsed = await asyncio.to_thread( self._renderer.parse_response, token_ids=completion.token_ids, + tools=self._tools, ) except Exception: logger.exception( diff --git a/torchtitan/experiments/rl/examples/alphabet_sort/config_registry.py b/torchtitan/experiments/rl/examples/alphabet_sort/config_registry.py index 26a9a72c8f..08454dc76f 100644 --- a/torchtitan/experiments/rl/examples/alphabet_sort/config_registry.py +++ b/torchtitan/experiments/rl/examples/alphabet_sort/config_registry.py @@ -13,6 +13,8 @@ import dataclasses +from renderers import GptOssRendererConfig, Qwen3RendererConfig + from torchtitan.components.checkpointer import CheckpointManager from torchtitan.components.loss import ChunkedLossWrapper from torchtitan.components.optimizer import default_adamw, LRSchedulersContainer @@ -23,6 +25,7 @@ ParallelismConfig, TrainingConfig, ) +from torchtitan.distributed.activation_checkpoint import FullAC from torchtitan.experiments.rl.actors.generator import ( SamplingConfig, VLLMCudagraphConfig, @@ -43,7 +46,7 @@ from torchtitan.experiments.rl.models.cast_linear import LMHeadCastConverter from torchtitan.experiments.rl.models.vllm_registry import InferenceParallelismConfig from torchtitan.experiments.rl.observability.metrics import MetricsProcessor -from torchtitan.experiments.rl.renderer import RendererConfig +from torchtitan.experiments.rl.renderer import RenderersLibraryConfig from torchtitan.experiments.rl.routing.inter_generator_router import ( InterGeneratorRouter, ) @@ -97,7 +100,9 @@ def rl_grpo_qwen3_0_6b_varlen() -> Controller.Config: ), compile=CompileConfig(enable=True, backend="aot_eager"), rollouter=AlphabetSortRollouter.Config(), - renderer=RendererConfig(name="qwen3", enable_thinking=False), + renderer=RenderersLibraryConfig( + renderers_config=Qwen3RendererConfig(enable_thinking=False) + ), generator_router=InterGeneratorRouter.Config( strategy=StickySessionRoutingStrategy.Config( fallback_strategy=LeastLoadedRoutingStrategy.Config() @@ -159,7 +164,9 @@ def rl_grpo_qwen3_0_6b_flex() -> Controller.Config: ), compile=CompileConfig(enable=True, backend="aot_eager"), rollouter=AlphabetSortRollouter.Config(), - renderer=RendererConfig(name="qwen3", enable_thinking=False), + renderer=RenderersLibraryConfig( + renderers_config=Qwen3RendererConfig(enable_thinking=False) + ), metrics=MetricsProcessor.Config(enable_wandb=True), trainer=PolicyTrainer.Config( optimizer=default_adamw(lr=2e-6), @@ -261,7 +268,9 @@ def rl_grpo_gpt_oss_20b_varlen() -> Controller.Config: ), compile=CompileConfig(enable=True, backend="aot_eager"), rollouter=AlphabetSortRollouter.Config(), - renderer=RendererConfig(name="gpt_oss", enable_thinking=False), + renderer=RenderersLibraryConfig( + renderers_config=GptOssRendererConfig(reasoning_effort="low") + ), generator_router=InterGeneratorRouter.Config( strategy=StickySessionRoutingStrategy.Config( fallback_strategy=LeastLoadedRoutingStrategy.Config() @@ -329,7 +338,9 @@ def rl_grpo_gpt_oss_debug_varlen() -> Controller.Config: # Debug tokenizer (vocab 2048, matches debugmodel); the gpt_oss renderer # needs gpt-oss special tokens absent here, so use the qwen3 renderer # like the other debug configs. - renderer=RendererConfig(name="qwen3", enable_thinking=False), + renderer=RenderersLibraryConfig( + renderers_config=Qwen3RendererConfig(enable_thinking=False) + ), metrics=MetricsProcessor.Config(enable_wandb=True), trainer=PolicyTrainer.Config( optimizer=default_adamw(lr=2e-6), @@ -397,7 +408,9 @@ def rl_grpo_gpt_oss_debug_varlen_batch_invariant() -> Controller.Config: # Debug tokenizer (vocab 2048, matches debugmodel); the gpt_oss renderer # needs gpt-oss special tokens absent here, so use the qwen3 renderer # like the other debug configs. - renderer=RendererConfig(name="qwen3", enable_thinking=False), + renderer=RenderersLibraryConfig( + renderers_config=Qwen3RendererConfig(enable_thinking=False) + ), metrics=MetricsProcessor.Config(enable_wandb=True), trainer=PolicyTrainer.Config( optimizer=default_adamw(lr=2e-6), @@ -457,7 +470,9 @@ def rl_grpo_qwen3_1_7b() -> Controller.Config: ), compile=CompileConfig(enable=True, backend="aot_eager"), rollouter=AlphabetSortRollouter.Config(), - renderer=RendererConfig(name="qwen3", enable_thinking=False), + renderer=RenderersLibraryConfig( + renderers_config=Qwen3RendererConfig(enable_thinking=False) + ), metrics=MetricsProcessor.Config(enable_wandb=True), trainer=PolicyTrainer.Config( optimizer=default_adamw(lr=2e-6), @@ -514,7 +529,9 @@ def rl_grpo_qwen3_14b() -> Controller.Config: ), compile=CompileConfig(enable=True, backend="aot_eager"), rollouter=AlphabetSortRollouter.Config(), - renderer=RendererConfig(name="qwen3", enable_thinking=False), + renderer=RenderersLibraryConfig( + renderers_config=Qwen3RendererConfig(enable_thinking=False) + ), metrics=MetricsProcessor.Config(enable_wandb=True), trainer=PolicyTrainer.Config( optimizer=default_adamw(lr=1e-6), @@ -582,7 +599,9 @@ def rl_grpo_qwen3_moe_debug_varlen() -> Controller.Config: # torch.compile and CUDA graph capture; disable both. compile=CompileConfig(enable=False), rollouter=AlphabetSortRollouter.Config(), - renderer=RendererConfig(name="qwen3", enable_thinking=False), + renderer=RenderersLibraryConfig( + renderers_config=Qwen3RendererConfig(enable_thinking=False) + ), metrics=MetricsProcessor.Config(enable_wandb=True), trainer=PolicyTrainer.Config( optimizer=default_adamw(lr=8e-4), @@ -713,7 +732,9 @@ def rl_grpo_qwen3_moe_debug_varlen_batch_invariant() -> Controller.Config: # torch.compile and CUDA graph capture; disable both. compile=CompileConfig(enable=False), rollouter=AlphabetSortRollouter.Config(), - renderer=RendererConfig(name="qwen3", enable_thinking=False), + renderer=RenderersLibraryConfig( + renderers_config=Qwen3RendererConfig(enable_thinking=False) + ), metrics=MetricsProcessor.Config(enable_wandb=True), trainer=PolicyTrainer.Config( optimizer=default_adamw(lr=8e-4), @@ -781,7 +802,9 @@ def rl_grpo_qwen3_30b_a3b_varlen() -> Controller.Config: ), compile=CompileConfig(enable=False), rollouter=AlphabetSortRollouter.Config(), - renderer=RendererConfig(name="qwen3", enable_thinking=False), + renderer=RenderersLibraryConfig( + renderers_config=Qwen3RendererConfig(enable_thinking=False) + ), metrics=MetricsProcessor.Config(enable_wandb=True), trainer=PolicyTrainer.Config( optimizer=default_adamw(lr=1e-6), @@ -892,7 +915,9 @@ def rl_grpo_qwen3_0_6b_varlen_batch_invariant() -> Controller.Config: ), compile=CompileConfig(enable=True, backend="aot_eager"), rollouter=AlphabetSortRollouter.Config(), - renderer=RendererConfig(name="qwen3", enable_thinking=False), + renderer=RenderersLibraryConfig( + renderers_config=Qwen3RendererConfig(enable_thinking=False) + ), metrics=MetricsProcessor.Config(enable_wandb=True), trainer=PolicyTrainer.Config( optimizer=default_adamw(lr=2e-6), @@ -971,7 +996,9 @@ def rl_grpo_qwen3_5_9b_varlen() -> Controller.Config: ), compile=CompileConfig(enable=False), rollouter=AlphabetSortRollouter.Config(), - renderer=RendererConfig(name="qwen3", enable_thinking=False), + renderer=RenderersLibraryConfig( + renderers_config=Qwen3RendererConfig(enable_thinking=False) + ), metrics=MetricsProcessor.Config(enable_wandb=True), trainer=PolicyTrainer.Config( optimizer=default_adamw(lr=1e-6), @@ -1055,7 +1082,9 @@ def rl_grpo_qwen3_5_debug_varlen() -> Controller.Config: ), compile=CompileConfig(enable=False), rollouter=AlphabetSortRollouter.Config(), - renderer=RendererConfig(name="qwen3", enable_thinking=False), + renderer=RenderersLibraryConfig( + renderers_config=Qwen3RendererConfig(enable_thinking=False) + ), metrics=MetricsProcessor.Config(enable_wandb=True), trainer=PolicyTrainer.Config( optimizer=default_adamw(lr=1e-6), @@ -1112,3 +1141,39 @@ def rl_grpo_qwen3_5_debug_varlen_batch_invariant() -> Controller.Config: config.generator, debug=_BATCH_INVARIANT_DEBUG ) return config + + +def rl_grpo_qwen3_6_27b_varlen_perf() -> Controller.Config: + """Qwen3.6-27B GRPO with fused OffsetRMSNorm on trainer and generator. + + Qwen3.6-27B uses the Qwen3.5-compatible dense Gated DeltaNet model flavor. + The 8-GPU layout assigns TP2 x FSDP2 to training and TP4 to generation. + """ + config = rl_grpo_qwen3_5_9b_varlen() + config.model_spec = _qwen3_5_rl_model_registry("27B", attn_backend="varlen") + config.hf_assets_path = "torchtitan/experiments/rl/example_checkpoint/Qwen3.6-27B" + perf_imports = ["torchtitan.overrides.offset_rmsnorm.triton_offset_rmsnorm"] + config.trainer = dataclasses.replace( + config.trainer, + optimizer=dataclasses.replace( + config.trainer.optimizer, + implementation="fused_opt_states_bf16", + ), + ac_config=FullAC.Config(), + parallelism=dataclasses.replace( + config.trainer.parallelism, + data_parallel_shard_degree=2, + tensor_parallel_degree=2, + ), + override=OverrideConfig(imports=list(perf_imports)), + ) + config.generator = dataclasses.replace( + config.generator, + parallelism=dataclasses.replace( + config.generator.parallelism, + data_parallel_degree=1, + tensor_parallel_degree=4, + ), + override=OverrideConfig(imports=list(perf_imports)), + ) + return config diff --git a/torchtitan/experiments/rl/examples/dapo_math/README.md b/torchtitan/experiments/rl/examples/dapo_math/README.md index 37affab39c..17bd6ec5dd 100644 --- a/torchtitan/experiments/rl/examples/dapo_math/README.md +++ b/torchtitan/experiments/rl/examples/dapo_math/README.md @@ -10,20 +10,20 @@ Each episode is single-turn: user math problem -> one assistant solution -> binary Math-Verify reward ``` -The prompt asks for step-by-step reasoning followed by a final `Answer:` expression. [Math-Verify](https://github.com/huggingface/Math-Verify) parses that expression and assigns a reward of one when it is mathematically equivalent to the reference answer, or zero otherwise. +The prompt asks for step-by-step reasoning followed by a final `Answer: \boxed{...}` expression. The scorer extracts the last boxed expression before [Math-Verify](https://github.com/huggingface/Math-Verify) checks it against the reference answer. An episode from the reference run is shown below. The prompt is reproduced in full; the response is abridged. ```text Prompt: Solve the following math problem step by step. The last line of your response -should be of the form Answer: $Answer (without quotes) where $Answer is the -answer to the problem. +should be of the form Answer: \boxed{$Answer}, where $Answer is the answer to +the problem. Let $r_1, r_2, \ldots, r_{47}$ be the roots of $x^{47} - 1 = 0$. Compute \( \sum_{i=1}^{47} r_i^{2020} \). -Remember to put your answer on its own line after "Answer:". +Remember to put your answer on its own line as "Answer: \boxed{...}". Response: The roots are the 47th roots of unity. Since 2020 is congruent to -1 diff --git a/torchtitan/experiments/rl/examples/dapo_math/config_registry.py b/torchtitan/experiments/rl/examples/dapo_math/config_registry.py index 6ccc2f8edd..4ccce8b6e7 100644 --- a/torchtitan/experiments/rl/examples/dapo_math/config_registry.py +++ b/torchtitan/experiments/rl/examples/dapo_math/config_registry.py @@ -8,6 +8,8 @@ from __future__ import annotations +from renderers import Qwen3RendererConfig + from torchtitan.components.checkpointer import CheckpointManager from torchtitan.components.loss import ChunkedLossWrapper from torchtitan.components.optimizer import default_adamw, LRSchedulersContainer @@ -33,7 +35,7 @@ from torchtitan.experiments.rl.models.cast_linear import LMHeadCastConverter from torchtitan.experiments.rl.models.vllm_registry import InferenceParallelismConfig from torchtitan.experiments.rl.observability.metrics import MetricsProcessor -from torchtitan.experiments.rl.renderer import RendererConfig +from torchtitan.experiments.rl.renderer import RenderersLibraryConfig from torchtitan.experiments.rl.routing.inter_generator_router import ( InterGeneratorRouter, ) @@ -81,7 +83,9 @@ def _qwen3_4b_dapo_math_config( ), ), ), - renderer=RendererConfig(name="qwen3", enable_thinking=True), + renderer=RenderersLibraryConfig( + renderers_config=Qwen3RendererConfig(enable_thinking=True) + ), num_generators=6, generator_router=InterGeneratorRouter.Config( strategy=LeastLoadedRoutingStrategy.Config() diff --git a/torchtitan/experiments/rl/examples/dapo_math/data.py b/torchtitan/experiments/rl/examples/dapo_math/data.py index 5146e908aa..0980fc37f9 100644 --- a/torchtitan/experiments/rl/examples/dapo_math/data.py +++ b/torchtitan/experiments/rl/examples/dapo_math/data.py @@ -14,13 +14,12 @@ from torchtitan.config import Configurable - -_AIME_PROMPT_TEMPLATE = ( +_MATH_PROMPT_TEMPLATE = ( "Solve the following math problem step by step. The last line of your response " - "should be of the form Answer: $Answer (without quotes) where $Answer is the " - "answer to the problem.\n\n" + "should be of the form Answer: \\boxed{{$Answer}}, where $Answer is the answer " + "to the problem.\n\n" "{problem}\n\n" - 'Remember to put your answer on its own line after "Answer:".' + 'Remember to put your answer on its own line as "Answer: \\boxed{{...}}".' ) @@ -102,7 +101,8 @@ def __init__(self, config: Config) -> None: raise ValueError("DAPO-Math rows must contain exactly one user prompt") samples.append( DapoMathSample( - prompt=prompt_messages[0]["content"], + # `prompt` is the raw question without answer-format instructions. + prompt=_MATH_PROMPT_TEMPLATE.format(problem=row["prompt"]), ground_truth=str(row["ground_truth"]), ) ) @@ -130,7 +130,7 @@ def __init__(self, config: Config) -> None: ).select(range(config.num_samples)) samples = [ DapoMathSample( - prompt=_AIME_PROMPT_TEMPLATE.format(problem=row["question"]), + prompt=_MATH_PROMPT_TEMPLATE.format(problem=row["question"]), ground_truth=str(row["answer"]), ) for row in dataset diff --git a/torchtitan/experiments/rl/examples/dapo_math/rubric.py b/torchtitan/experiments/rl/examples/dapo_math/rubric.py index 2b76f9d623..da8076c3ce 100644 --- a/torchtitan/experiments/rl/examples/dapo_math/rubric.py +++ b/torchtitan/experiments/rl/examples/dapo_math/rubric.py @@ -8,44 +8,50 @@ from dataclasses import dataclass -from math_verify import LatexExtractionConfig, LatexNormalizationConfig, parse, verify +from math_verify import parse, verify from math_verify.errors import TimeoutException from torchtitan.experiments.rl.examples.dapo_math.data import DapoMathSample from torchtitan.experiments.rl.rollout import Rollout from torchtitan.experiments.rl.rubrics import RewardFn +_BOXED_START = r"\boxed{" -# Require an `Answer:` or `\boxed{}` marker so intermediate math is ignored. -_FINAL_ANSWER_EXTRACTION = [ - LatexExtractionConfig( - normalization_config=LatexNormalizationConfig(units=True), - boxed_match_priority=0, - try_extract_without_anchor=False, - ) -] + +def _last_boxed_expression(text: str) -> str | None: + """Return the last complete `\\boxed{...}` expression.""" + start = text.rfind(_BOXED_START) + if start == -1: + return None + + answer_start = start + len(_BOXED_START) + depth = 1 + for index, char in enumerate(text[answer_start:], start=answer_start): + depth += (char == "{") - (char == "}") + if depth == 0: + return text[start : index + 1] + return None def score_math_response(response: str, ground_truth: str) -> float: - """Score an `Answer:` or `\\boxed{}` expression with Math-Verify. + """Score the final `\\boxed{}` expression with Math-Verify. Args: - response: Model response containing a marked final answer. + response: Model response containing a boxed final answer. ground_truth: Expected answer from the dataset. Example: - score_math_response("work\nAnswer: $34$", "34") # 1.0 + score_math_response(r"work\nAnswer: \boxed{34}", "34") # 1.0 """ + prediction = _last_boxed_expression(response) + if prediction is None: + return 0.0 + try: # TODO: Re-enable Math-Verify timeouts after resolving its signal-based # timeout failure in rollout worker threads (signals require the main thread). gold = parse(ground_truth, parsing_timeout=None) - prediction = parse( - response, - extraction_config=_FINAL_ANSWER_EXTRACTION, - extraction_mode="first_match", - parsing_timeout=None, - ) + prediction = parse(prediction, parsing_timeout=None) return float(bool(gold) and verify(gold, prediction, timeout_seconds=None)) except (Exception, TimeoutException): # Model output is untrusted; malformed LaTeX is an incorrect answer, not a diff --git a/torchtitan/experiments/rl/examples/search_r1/config_registry.py b/torchtitan/experiments/rl/examples/search_r1/config_registry.py index 6f71e4ddf2..cc6153b26a 100644 --- a/torchtitan/experiments/rl/examples/search_r1/config_registry.py +++ b/torchtitan/experiments/rl/examples/search_r1/config_registry.py @@ -18,6 +18,8 @@ import dataclasses +from renderers import Qwen3RendererConfig + from torchtitan.components.checkpointer import CheckpointManager from torchtitan.components.loss import ChunkedLossWrapper from torchtitan.components.optimizer import default_adamw, LRSchedulersContainer @@ -44,9 +46,12 @@ SearchR1Worker, ) from torchtitan.experiments.rl.losses import DAPOLoss +from torchtitan.experiments.rl.models.muse_glimmer.renderer import ( + MuseGlimmerRendererConfig, +) from torchtitan.experiments.rl.models.vllm_registry import InferenceParallelismConfig from torchtitan.experiments.rl.observability.metrics import MetricsProcessor -from torchtitan.experiments.rl.renderer import RendererConfig +from torchtitan.experiments.rl.renderer import RenderersLibraryConfig 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 ( @@ -78,7 +83,9 @@ def rl_grpo_qwen3_1_7b_search_r1() -> Controller.Config: advantage=AdvantageEstimator.Config(should_std_normalize=True), ), ), - renderer=RendererConfig(name="qwen3", enable_thinking=False), + renderer=RenderersLibraryConfig( + renderers_config=Qwen3RendererConfig(enable_thinking=False) + ), metrics=MetricsProcessor.Config(enable_wandb=True), trainer=PolicyTrainer.Config( optimizer=default_adamw(lr=1e-6), @@ -206,7 +213,9 @@ def rl_grpo_qwen3_30b_a3b_deepep_search_r1_perf() -> Controller.Config: advantage=AdvantageEstimator.Config(should_std_normalize=True), ), ), - renderer=RendererConfig(name="qwen3", enable_thinking=False), # TODO: TBD + renderer=RenderersLibraryConfig( + renderers_config=Qwen3RendererConfig(enable_thinking=False) + ), metrics=MetricsProcessor.Config(enable_wandb=True), trainer=PolicyTrainer.Config( optimizer=default_adamw(lr=1e-6), @@ -280,12 +289,9 @@ def rl_grpo_muse_glimmer_30b_search_r1() -> Controller.Config: 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 + on load, and the renderer 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 @@ -306,7 +312,7 @@ def rl_grpo_muse_glimmer_30b_search_r1() -> Controller.Config: advantage=AdvantageEstimator.Config(should_std_normalize=True), ), ), - renderer=RendererConfig(name="muse_glimmer", enable_thinking=True), + renderer=MuseGlimmerRendererConfig(), metrics=MetricsProcessor.Config(enable_wandb=True), trainer=PolicyTrainer.Config( optimizer=default_adamw(lr=1e-6), diff --git a/torchtitan/experiments/rl/generate.py b/torchtitan/experiments/rl/generate.py index 5b3f5ab45f..805fdeb382 100755 --- a/torchtitan/experiments/rl/generate.py +++ b/torchtitan/experiments/rl/generate.py @@ -187,7 +187,8 @@ def generate() -> None: logger.debug("vLLM LLMEngine initialized successfully") - renderer = config.renderer.build(tokenizer_path=model_path) + tokenizer = config.tokenizer.build(tokenizer_path=model_path) + renderer = config.renderer.build(tokenizer=tokenizer) stop_token_ids = list(renderer.get_stop_token_ids()) # Create sampling parameters from config diff --git a/torchtitan/experiments/rl/models/muse_glimmer/renderer.py b/torchtitan/experiments/rl/models/muse_glimmer/renderer.py index 2be5350f70..f0b7c4208b 100644 --- a/torchtitan/experiments/rl/models/muse_glimmer/renderer.py +++ b/torchtitan/experiments/rl/models/muse_glimmer/renderer.py @@ -28,18 +28,13 @@ reasoning states. Treat that test as the spec -- if the template changes upstream, it fails first. -Implements the ``renderers.Renderer`` Protocol. ``register()`` installs it into the -``renderers`` library's public registry (``RENDERER_REGISTRY`` / ``_CONFIG_BY_NAME``), -which is that library's supported extension path -- no fork or upstream change needed. - -Every other TorchTitan model resolves to a renderer that lives in -PrimeIntellect-ai/renderers. This one ships here because Muse Glimmer is not in that -library yet. +Implements the ``renderers.Renderer`` Protocol. Muse Glimmer is not in the library yet, so +``MuseGlimmerRendererConfig`` is a TorchTitan ``RendererConfig`` whose ``build`` constructs +this class directly. TODO: upstream this to PrimeIntellect-ai/renderers (renderer -> renderers/muse_glimmer.py, -atem.py -> a tool parser in renderers/parsers.py), then delete both files and -``register()``, leaving only the ``_RENDERER_BY_MODEL`` entry in -experiments/rl/renderer.py. +atem.py -> a tool parser in renderers/parsers.py), then delete both files and select it +through ``RenderersLibraryConfig`` like the other renderers. It lives under ``experiments/rl`` rather than ``torchtitan/models/muse_glimmer`` because RL is its only consumer and ``renderers`` is an RL-only optional dependency; keeping it @@ -51,8 +46,10 @@ import datetime import json import re -from typing import ClassVar, Literal, NamedTuple +from dataclasses import dataclass, replace +from typing import NamedTuple +from renderers import Renderer from renderers.base import ( extract_message_tool_names, ParsedResponse, @@ -63,31 +60,18 @@ should_rerender_for_thinking_retention, trim_to_turn_close, ) -from renderers.configs import BaseRendererConfig +from renderers.configs import ThinkingRetention -from .atem import parse_atem_tool_calls, render_atem_tool_call +from torchtitan.components.tokenizer import HuggingFaceTokenizer +from torchtitan.experiments.rl.renderer import RendererConfig, RendererTokenizerWrapper -RENDERER_NAME = "muse_glimmer" +from .atem import parse_atem_tool_calls, render_atem_tool_call -class MuseGlimmerRendererConfig(BaseRendererConfig): +@dataclass(kw_only=True, slots=True) +class MuseGlimmerRendererConfig(RendererConfig): """Muse Glimmer (harmony chat format + ATEM tool calls) renderer config.""" - name: Literal["muse_glimmer"] = RENDERER_NAME - - # renderers validates in BaseRendererConfig.__pydantic_init_subclass__ that every - # non-base field is classified as either a chat-template kwarg or a renderer-internal - # knob; the two sets must be disjoint and together cover all of them. Declared - # unconditionally -- versions without the validator ignore these ClassVars, so this - # is compatible with both. The template fields mirror kwargs the published - # chat_template.jinja reads, which is what the library's parity matrix varies. - _template_fields: ClassVar[frozenset[str]] = frozenset( - {"reasoning_strength", "knowledge_cutoff", "current_date"} - ) - _internal_fields: ClassVar[frozenset[str]] = frozenset( - {"retain_reasoning_in_history", "answer_from_reasoning_fallback"} - ) - reasoning_strength: str | None = None """Sizes the reasoning budget, rendered as ``Reasoning strength: .`` @@ -131,6 +115,34 @@ class MuseGlimmerRendererConfig(BaseRendererConfig): empty ``content`` is unscoreable and the answer is often the final reasoning line. """ + thinking_retention: ThinkingRetention | None = None + """The library-wide bridge policy override (`renderers.BaseRendererConfig.thinking_retention`). + ``None`` keeps the template's implied policy; ``"tool_cycle"`` re-renders at a new user query.""" + + def __post_init__(self) -> None: + # A dataclass does not validate values; these knobs change bridging and reward + # scoring with nothing visible in the rendered prompt, so check them here. + for name in ("reasoning_strength", "knowledge_cutoff", "current_date"): + value = getattr(self, name) + if value is not None and not isinstance(value, str): + raise TypeError( + f"{name} must be str | None, got {type(value).__name__}" + ) + for name in ("retain_reasoning_in_history", "answer_from_reasoning_fallback"): + value = getattr(self, name) + if type(value) is not bool: + raise TypeError(f"{name} must be bool, got {type(value).__name__}") + if self.thinking_retention not in (None, "tool_cycle", "all"): + raise ValueError( + "thinking_retention must be None, 'tool_cycle' or 'all', " + f"got {self.thinking_retention!r}" + ) + + def build(self, *, tokenizer: HuggingFaceTokenizer) -> Renderer: + # Snapshot the config, as `Configurable.Config.build` does, so later edits to the + # recipe object cannot desynchronize full renders from bridging. + return MuseGlimmerRenderer(RendererTokenizerWrapper(tokenizer), replace(self)) + # Muse Glimmer special tokens. The ids are checked against the tokenizer in __init__ # rather than trusted, since they are baked into parse_response and the loss mask. @@ -343,17 +355,13 @@ class _Piece(NamedTuple): class MuseGlimmerRenderer: def __init__(self, tokenizer, config: MuseGlimmerRendererConfig | None = None): - # (tokenizer, config) is the renderers-library constructor contract, so - # ``create_renderer`` can instantiate this from RENDERER_REGISTRY. + # Match the `(tokenizer, config)` constructor used by library renderers. self._tok = tokenizer self._config = config or MuseGlimmerRendererConfig() - # The controller reads renderer._tokenizer (e.g. for pad_id=eos_token_id). - self._tokenizer = tokenizer self._bos = tokenizer.bos_token or "" - # BaseRendererConfig.thinking_retention is the library-wide knob every renderer - # is expected to honour in its bridge. Muse Glimmer's published chat template - # renders reasoning_content for every assistant turn unconditionally -- no - # query-boundary drop like gpt-oss's auto_drop_analysis or Qwen3's think-block + # `thinking_retention` is the library-wide bridge knob. Muse Glimmer's published + # chat template renders reasoning_content for every assistant turn unconditionally + # -- no query-boundary drop like gpt-oss's auto_drop_analysis or Qwen3's think-block # stripping -- so "all" is the template-faithful implied policy. An explicit # thinking_retention on the config overrides it. self.effective_thinking_retention = resolve_thinking_retention( @@ -758,31 +766,3 @@ def parse_response(self, token_ids, *, tools=None) -> ParsedResponse: reasoning_content=reasoning, tool_calls=tool_calls, ) - - -def register() -> None: - """Install the muse_glimmer renderer into the ``renderers`` library registry. - - Uses the library's public extension surface -- implement the ``Renderer`` - Protocol, then add the class to ``RENDERER_REGISTRY`` and its config to - ``_CONFIG_BY_NAME`` -- so ``create_renderer(config_from_name("muse_glimmer"))`` - resolves it. Also maps the ``muse_glimmer`` TorchTitan model name to it, which is what - ``RendererConfig(name="muse_glimmer")`` looks up. - - Idempotent. Delete this once the renderer is upstreamed to - PrimeIntellect-ai/renderers (only the _RENDERER_BY_MODEL entry stays). - """ - from renderers import base as renderers_base, configs as renderers_configs - - from torchtitan.experiments.rl.renderer import _RENDERER_BY_MODEL - - # Populate the library's built-ins first: _populate_registry() early-returns if - # RENDERER_REGISTRY is already non-empty, so registering before it runs would - # suppress every built-in renderer. - renderers_base._populate_registry() - - renderers_configs._CONFIG_BY_NAME.setdefault( - RENDERER_NAME, MuseGlimmerRendererConfig - ) - renderers_base.RENDERER_REGISTRY[RENDERER_NAME] = MuseGlimmerRenderer - _RENDERER_BY_MODEL["muse_glimmer"] = RENDERER_NAME diff --git a/torchtitan/experiments/rl/renderer.py b/torchtitan/experiments/rl/renderer.py index 013f945ae7..76279fa282 100644 --- a/torchtitan/experiments/rl/renderer.py +++ b/torchtitan/experiments/rl/renderer.py @@ -6,109 +6,142 @@ from __future__ import annotations -import logging -from dataclasses import dataclass, fields +from dataclasses import dataclass +from typing import Annotated, Any -from renderers import config_from_name, create_renderer, Renderer +import tyro +from renderers import create_renderer, Renderer +from renderers.configs import BaseRendererConfig +from torchtitan.components.tokenizer import HuggingFaceTokenizer from torchtitan.config import Configurable -logger = logging.getLogger(__name__) - -# Map a TorchTitan model name to its `renderers` renderer. Models not listed fall -# back to "auto" (renderers resolves from the tokenizer) -# https://github.com/PrimeIntellect-ai/renderers/blob/942449c37ab6e9fab26d59b40336514c8baa6b13/renderers/configs.py#L404 -_RENDERER_BY_MODEL = { - "qwen3": "qwen3", - "qwen3_vl": "qwen3-vl", - "gpt_oss": "gpt-oss", - "deepseek_v3": "deepseek-v3", - # TODO: upstream the Muse Glimmer renderer to PrimeIntellect-ai/renderers, then - # delete its `register()` and point this at the library's name (hyphenated, like - # the entries above). It ships in torchtitan and self-registers only because the - # library has no Muse Glimmer renderer yet; every other model here resolves to one - # the library owns. See rl/models/muse_glimmer/renderer.py. - "muse_glimmer": "muse_glimmer", - "default": "default", # llama3 - "auto": "auto", # ignores knobs, resolves from tokenizer, -} - @dataclass(kw_only=True, slots=True) class RendererConfig(Configurable.Config): - """Selects the renderer used for chat message <-> token conversion. + """Base config of a renderer; `build` returns a `renderers.Renderer` on TorchTitan's tokenizer. + + Subclasses: `RenderersLibraryConfig` for a renderer from the `renderers` library, and + in-tree renderers such as `MuseGlimmerRendererConfig`. + """ - Wraps `PrimeIntellect-ai/renderers`. `build` loads a tokenizer from - `tokenizer_path`, maps the model `name` to a renderer, and forwards any - supported knobs. + def build(self, *, tokenizer: HuggingFaceTokenizer) -> Renderer: + raise NotImplementedError - Args: - name: TorchTitan model name (e.g. `"qwen3"`, `"llama3"`), mapped to a - `renderers` renderer via `_RENDERER_BY_MODEL`. `None` (the default) - resolves the renderer from the tokenizer. - tool_parser: Tool-call parser name, when the renderer supports it. - reasoning_parser: Reasoning parser name, when the renderer supports it. - enable_thinking: Let the model emit reasoning, when supported. - preserve_all_thinking: Keep historical reasoning in future prompts. - preserve_thinking_between_tool_calls: Keep reasoning during tool loops. - Every field defaults to `None`; a non-`None` value overrides that knob on the - chosen renderer's config, otherwise the renderer keeps its own default. +@dataclass(kw_only=True, slots=True) +class RenderersLibraryConfig(RendererConfig): + """Builds one of the `renderers` library's renderers on TorchTitan's tokenizer. Example: - renderer = RendererConfig(name="qwen3").build(tokenizer_path="./Qwen3-0.6B") + from renderers import Qwen3RendererConfig + + from torchtitan.components.tokenizer import HuggingFaceTokenizer + from torchtitan.experiments.rl.renderer import RenderersLibraryConfig + + renderer = RenderersLibraryConfig( + renderers_config=Qwen3RendererConfig(enable_thinking=False) + ).build(tokenizer=HuggingFaceTokenizer(tokenizer_path="./Qwen3-0.6B")) prompt_ids = renderer.render_ids( - [{"role": "user", "content": "hi"}], add_generation_prompt=True + [{"role": "user", "content": "hi"}], + add_generation_prompt=True, ) """ - name: str | None = None - tool_parser: str | None = None - reasoning_parser: str | None = None - enable_thinking: bool | None = None - preserve_all_thinking: bool | None = None - preserve_thinking_between_tool_calls: bool | None = None - - def build(self, *, tokenizer_path: str) -> Renderer: - # TODO(renderers#70): use TorchTitan's tokenizer once `renderers` supports - # bring-your-own-tokenizer (PR adds a Tokenizer protocol; drops transformers). - from transformers import AutoTokenizer - - tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) - - # `name=None` (or "auto") -> let `create_renderer` resolve from the tokenizer. - renderer_name = _RENDERER_BY_MODEL.get(self.name, self.name) - if renderer_name == "muse_glimmer": - # TODO: temporary. Delete this block once the Muse Glimmer renderer is - # upstreamed to PrimeIntellect-ai/renderers -- the library registers its - # own renderers in _populate_registry(), so no torchtitan-side hook is - # needed for any other model here. - # - # Until then it has to live in build(), not in the config registry: the - # renderer is constructed inside a Monarch-spawned RolloutWorker, which - # only receives the serialized config and never imports the recipe module, - # so a register() call there never runs in that process. - from torchtitan.experiments.rl.models.muse_glimmer import ( - renderer as _muse_glimmer_renderer, + renderers_config: Annotated[BaseRendererConfig, tyro.conf.Suppress] + """The library's typed config for the model, e.g. `Qwen3RendererConfig(enable_thinking=False)`. + Renderers and their options: + https://github.com/PrimeIntellect-ai/renderers/blob/renderers-v0.1.11/docs/renderer-config.md""" + + def to_dict(self) -> dict[str, Any]: + return {"renderers_config": self.renderers_config.model_dump(mode="json")} + + def build(self, *, tokenizer: HuggingFaceTokenizer) -> Renderer: + if self.renderers_config.name == "auto": + raise ValueError( + f"AutoRendererConfig resolves by exact match of tokenizer.name_or_path ({tokenizer.tokenizer_path!r}) " + "against renderers' MODEL_RENDERER_MAP, else falls back to DefaultRenderer (unsupported here). " + "Pick the model's renderer, e.g. Qwen3RendererConfig(...)." + ) + if self.renderers_config.name == "default": + raise ValueError( + "DefaultRenderer needs Hugging Face apply_chat_template; TorchTitan's template rendering lacks " + "its special-token variables (bos_token, ...) and would silently produce different tokens. " + "Pick the model's renderer, e.g. Qwen3RendererConfig(...)." ) + return create_renderer( + tokenizer=RendererTokenizerWrapper(tokenizer), config=self.renderers_config + ) + + +class RendererTokenizerWrapper: + """Adapt TorchTitan's loaded tokenizer to `renderers.OffsetTokenizer`. + + Protocol and bring-your-own-tokenizer guide: + https://github.com/PrimeIntellect-ai/renderers/blob/renderers-v0.1.11/renderers/base.py#L668-L699 + https://github.com/PrimeIntellect-ai/renderers/blob/renderers-v0.1.11/README.md#install + + `renderers` needs Hugging Face-style special-token attributes, raw encoding + without automatic BOS/EOS, token-to-id lookup, and character offsets. The + offsets identify tokens from message content (`is_content`). This adapter + exposes that interface from TorchTitan's underlying `tokenizers.Tokenizer`; + it does not load a second tokenizer. + + Example: + + from torchtitan.components.tokenizer import HuggingFaceTokenizer + from torchtitan.experiments.rl.renderer import RendererTokenizerWrapper + + tokenizer = RendererTokenizerWrapper( + HuggingFaceTokenizer(tokenizer_path="./Qwen3-0.6B") + ) + tokenizer.encode("hi") # [6023] + tokenizer( + "hi", + add_special_tokens=False, + return_offsets_mapping=True, + ) # ids + character offsets + tokenizer.convert_tokens_to_ids("<|im_end|>") # 151645 + """ + + def __init__(self, tokenizer: HuggingFaceTokenizer): + # The `tokenizers.Tokenizer` inside; it has the offsets and token -> id lookup. + self._tokenizer_backend = tokenizer.tokenizer + self.name_or_path = tokenizer.tokenizer_path + self.bos_token = tokenizer.bos_token + self.eos_token = tokenizer.eos_token + self.bos_token_id = tokenizer.bos_id + self.eos_token_id = tokenizer.eos_id + # `tokenizers` returns None for unknown tokens; it has no unk id. + self.unk_token_id = None + + def encode( + self, text: str, add_special_tokens: bool = False, **kwargs + ) -> list[int]: + return self._tokenizer_backend.encode( + text, add_special_tokens=add_special_tokens + ).ids + + def decode(self, token_ids, skip_special_tokens: bool = False, **kwargs) -> str: + return self._tokenizer_backend.decode( + list(token_ids), skip_special_tokens=skip_special_tokens + ) - _muse_glimmer_renderer.register() - renderer_config = config_from_name(renderer_name) if renderer_name else None - if renderer_config is None: - return create_renderer(tokenizer, None) - - # Rebuild the typed config and pass parameters - # that are not None and are supported - config_type = type(renderer_config) - args = { - field.name: getattr(self, field.name) # {key: value} - for field in fields(self) - if field.name != "name" # Get all self.fields, except name - and getattr(self, field.name) is not None # Only consider provided fields - and field.name in config_type.model_fields # Config supports this field - } - logger.info( - f"Using renderer {renderer_name}, of type {config_type}, with args {args}" + def convert_tokens_to_ids( + self, tokens: str | list[str] + ) -> int | None | list[int | None]: + if isinstance(tokens, str): + return self._tokenizer_backend.token_to_id(tokens) + return [self._tokenizer_backend.token_to_id(token) for token in tokens] + + def __call__( + self, text: str, *, add_special_tokens: bool, return_offsets_mapping: bool + ) -> dict: + encoding = self._tokenizer_backend.encode( + text, add_special_tokens=add_special_tokens ) - return create_renderer(tokenizer, config_type(**args)) + output = {"input_ids": encoding.ids} + if return_offsets_mapping: + output["offset_mapping"] = encoding.offsets + return output diff --git a/torchtitan/experiments/rl/requirements.txt b/torchtitan/experiments/rl/requirements.txt index 9ef9658732..46b8ca477c 100644 --- a/torchtitan/experiments/rl/requirements.txt +++ b/torchtitan/experiments/rl/requirements.txt @@ -3,4 +3,4 @@ opentelemetry-sdk opentelemetry-exporter-otlp-proto-http pygtrie portpicker -git+https://github.com/PrimeIntellect-ai/renderers.git@main +renderers==0.1.11 diff --git a/torchtitan/experiments/rl/rollout/rollouter.py b/torchtitan/experiments/rl/rollout/rollouter.py index 44e7be7c19..1f468284c7 100644 --- a/torchtitan/experiments/rl/rollout/rollouter.py +++ b/torchtitan/experiments/rl/rollout/rollouter.py @@ -13,8 +13,10 @@ from monarch.actor import ProcMesh, this_host +from torchtitan.components.tokenizer import HuggingFaceTokenizer from torchtitan.config import Configurable from torchtitan.experiments.rl.environment import MessageEnv, TokenEnv +from torchtitan.experiments.rl.renderer import RendererConfig from torchtitan.experiments.rl.rollout.advantage import AdvantageEstimator from torchtitan.experiments.rl.rollout.types import ( GenerateFn, @@ -33,7 +35,6 @@ from torchtitan.experiments.rl.actors.generator import SamplingConfig from torchtitan.experiments.rl.actors.rollout_worker import RolloutWorkerActor - from torchtitan.experiments.rl.renderer import RendererConfig logger = logging.getLogger(__name__) @@ -135,6 +136,7 @@ def get_validation_sample(self) -> object: async def setup_async( self, *, + tokenizer_config: HuggingFaceTokenizer.Config, renderer_config: RendererConfig, hf_assets_path: str, ) -> None: @@ -155,6 +157,7 @@ async def setup_async( num_threads=self._config.num_threads_per_worker, ) await self._worker_actors.setup_async.call( + tokenizer_config=tokenizer_config, renderer_config=renderer_config, hf_assets_path=hf_assets_path, ) @@ -241,11 +244,13 @@ def __init__(self, config: Config) -> None: async def setup_async( self, *, + tokenizer_config: HuggingFaceTokenizer.Config, renderer_config: RendererConfig, hf_assets_path: str, ) -> None: """Build runtime dependencies after the worker actor is spawned.""" - self._renderer = renderer_config.build(tokenizer_path=hf_assets_path) + tokenizer = tokenizer_config.build(tokenizer_path=hf_assets_path) + self._renderer = renderer_config.build(tokenizer=tokenizer) def make_env_group( self, diff --git a/torchtitan/experiments/rl/tests/integration_tests.py b/torchtitan/experiments/rl/tests/integration_tests.py index 11a1d5f92e..cdf429005a 100644 --- a/torchtitan/experiments/rl/tests/integration_tests.py +++ b/torchtitan/experiments/rl/tests/integration_tests.py @@ -46,7 +46,6 @@ def build_rl_test_list() -> list[OverrideDefinitions]: "--async-loop.num-samples-per-prompt 2", "--trainer.training.max_context_length 1024", "--trainer.training.num_tokens_per_microbatch_per_dp_rank 2048", - "--renderer.enable-thinking False", "--generator.sampling.max_tokens 256", "--trainer.debug.no_batch_invariant", "--generator.debug.no_batch_invariant", @@ -75,7 +74,6 @@ def build_rl_test_list() -> list[OverrideDefinitions]: "--async-loop.num-samples-per-prompt 2", "--trainer.training.max_context_length 1024", "--trainer.training.num_tokens_per_microbatch_per_dp_rank 2048", - "--renderer.enable-thinking False", "--generator.sampling.max_tokens 256", "--trainer.debug.no_batch_invariant", "--generator.debug.no_batch_invariant", @@ -104,7 +102,6 @@ def build_rl_test_list() -> list[OverrideDefinitions]: "--async-loop.num-samples-per-prompt 2", "--trainer.training.max_context_length 1024", "--trainer.training.num_tokens_per_microbatch_per_dp_rank 2048", - "--renderer.enable-thinking False", "--generator.sampling.max_tokens 256", "--trainer.debug.no_batch_invariant", "--generator.debug.no_batch_invariant", @@ -145,7 +142,6 @@ def build_rl_test_list() -> list[OverrideDefinitions]: "--async-loop.num-samples-per-prompt 2", "--trainer.training.max_context_length 1024", "--trainer.training.num_tokens_per_microbatch_per_dp_rank 2048", - "--renderer.enable-thinking False", "--generator.sampling.max_tokens 256", "--trainer.debug.no_batch_invariant", "--generator.debug.no_batch_invariant", @@ -166,7 +162,6 @@ def build_rl_test_list() -> list[OverrideDefinitions]: "--async-loop.num-samples-per-prompt 2", "--trainer.training.max_context_length 1024", "--trainer.training.num_tokens_per_microbatch_per_dp_rank 2048", - "--renderer.enable-thinking False", "--generator.sampling.max_tokens 256", "--trainer.debug.no_batch_invariant", "--generator.debug.no_batch_invariant", @@ -199,7 +194,6 @@ def build_rl_test_list() -> list[OverrideDefinitions]: "--async-loop.num-samples-per-prompt 2", "--trainer.training.max_context_length 1024", "--trainer.training.num_tokens_per_microbatch_per_dp_rank 2048", - "--renderer.enable-thinking False", "--generator.sampling.max_tokens 128", "--metrics.no-enable-wandb", ], @@ -220,7 +214,6 @@ def build_rl_test_list() -> list[OverrideDefinitions]: "--async-loop.num-samples-per-prompt 2", "--trainer.training.max_context_length 1024", "--trainer.training.num_tokens_per_microbatch_per_dp_rank 2048", - "--renderer.enable-thinking False", "--generator.sampling.max_tokens 256", "--trainer.checkpoint.no-enable", # use random-init weights "--generator.checkpoint.no-enable", @@ -247,7 +240,6 @@ def build_rl_test_list() -> list[OverrideDefinitions]: "--async-loop.num-samples-per-prompt 2", "--trainer.training.max_context_length 1024", "--trainer.training.num_tokens_per_microbatch_per_dp_rank 1024", - "--renderer.enable-thinking False", "--generator.sampling.max_tokens 128", "--trainer.checkpoint.no-enable", # random-init weights "--generator.checkpoint.no-enable", diff --git a/torchtitan/experiments/rl/tests/test_alphabet_sort.py b/torchtitan/experiments/rl/tests/test_alphabet_sort.py index d9665a9535..8d4305e93b 100644 --- a/torchtitan/experiments/rl/tests/test_alphabet_sort.py +++ b/torchtitan/experiments/rl/tests/test_alphabet_sort.py @@ -11,6 +11,9 @@ import asyncio import pytest +from renderers import Qwen3RendererConfig + +from torchtitan.components.tokenizer import HuggingFaceTokenizer from torchtitan.experiments.rl.examples.alphabet_sort import ( AlphabetSortDataset, @@ -22,6 +25,7 @@ ) from torchtitan.experiments.rl.examples.alphabet_sort.env import AlphabetSortEnv from torchtitan.experiments.rl.examples.alphabet_sort.rubric import score_sorted_list +from torchtitan.experiments.rl.renderer import RenderersLibraryConfig from torchtitan.experiments.rl.rollout import Rollout, RolloutStatus, RolloutTurn from torchtitan.experiments.rl.types import RolloutTurnID @@ -369,10 +373,6 @@ def test_env_walks_through_follow_up_turns() -> None: def test_rollouter_builds_one_env_per_group_member( monkeypatch: pytest.MonkeyPatch, ) -> None: - class _RendererConfig: - def build(self, *, tokenizer_path: str): - return None - _patch_names(monkeypatch) config = AlphabetSortRollouter.Config() rollouter = AlphabetSortRollouter(config) @@ -381,8 +381,11 @@ def build(self, *, tokenizer_path: str): assert isinstance(worker, AlphabetSortWorker) asyncio.run( worker.setup_async( - renderer_config=_RendererConfig(), - hf_assets_path="hf_assets_path", + tokenizer_config=HuggingFaceTokenizer.Config(), + renderer_config=RenderersLibraryConfig( + renderers_config=Qwen3RendererConfig(enable_thinking=False) + ), + hf_assets_path="tests/assets/tokenizer", ) ) sample = rollouter.get_training_sample() diff --git a/torchtitan/experiments/rl/tests/test_dapo_math.py b/torchtitan/experiments/rl/tests/test_dapo_math.py index a5a106bbda..7427124cdc 100644 --- a/torchtitan/experiments/rl/tests/test_dapo_math.py +++ b/torchtitan/experiments/rl/tests/test_dapo_math.py @@ -18,10 +18,12 @@ DapoMathDataset, DapoMathEnv, DapoMathSample, - data as math_data, RewardMathVerify, score_math_response, ) +from torchtitan.experiments.rl.examples.dapo_math import ( + data as math_data, +) from torchtitan.experiments.rl.rollout import Rollout, RolloutStatus, RolloutTurn from torchtitan.experiments.rl.types import RolloutTurnID @@ -30,14 +32,17 @@ def _dapo_rows() -> list[dict]: return [ { "source_prompt": [{"role": "user", "content": "problem 1"}], + "prompt": "problem 1", "ground_truth": "34", }, { "source_prompt": [{"role": "user", "content": "problem 2"}], + "prompt": "problem 2", "ground_truth": "113", }, { "source_prompt": [{"role": "user", "content": "problem 3"}], + "prompt": "problem 3", "ground_truth": "7", }, ] @@ -55,6 +60,7 @@ def test_dapo_dataset_is_deterministic_and_resumable(monkeypatch) -> None: resumed = config.build() resumed.load_state_dict(checkpoint) assert [next(resumed) for _ in range(3)] == expected + assert all(r"Answer: \boxed{" in sample.prompt for sample in expected) def test_aime_dataset_combines_both_subsets(monkeypatch) -> None: @@ -69,6 +75,7 @@ def load_dataset(repo_id, subset, *, split): assert [sample.ground_truth for sample in samples] == [r"42^\circ", r"\boxed{42}"] assert "AIME2025-I question" in samples[0].prompt assert "AIME2025-II question" in samples[1].prompt + assert all(r"Answer: \boxed{" in sample.prompt for sample in samples) def test_aime_dataset_restarts_after_configured_num_samples(monkeypatch) -> None: @@ -108,20 +115,32 @@ def _rollout(response: str) -> Rollout: ) -def test_math_verifier_requires_an_answer_marker() -> None: - assert score_math_response("work\nAnswer: $34$", "34") == 1.0 - assert score_math_response(r"work\n\boxed{34}", "34") == 1.0 +def test_math_verifier_requires_a_boxed_answer() -> None: + assert score_math_response(r"work\nAnswer: \boxed{34}", "34") == 1.0 + assert score_math_response(r"work\n\boxed{\frac{68}{2}}", "34") == 1.0 + assert score_math_response("work\nAnswer: $34$", "34") == 0.0 + assert score_math_response("work\nAnswer: 34", "34") == 0.0 assert score_math_response("work mentions 34", "34") == 0.0 +def test_math_verifier_uses_the_last_boxed_answer() -> None: + response = r"Work: \boxed{2003^{2002^{2001}}}" "\n" r"Answer: \boxed{34}" + assert score_math_response(response, "34") == 1.0 + + +def test_math_verifier_rejects_unboxed_large_intermediate_expression() -> None: + response = r"Work: \[2003^{2002^{2001}}\]" "\n" r"Final: \[Answer: 009\]" + assert score_math_response(response, "241") == 0.0 + + def test_math_verifier_works_in_rollout_worker_thread() -> None: with ThreadPoolExecutor(max_workers=1) as executor: - result = executor.submit(score_math_response, "work\nAnswer: $34$", "34") + result = executor.submit(score_math_response, r"work\nAnswer: \boxed{34}", "34") assert result.result() == 1.0 def test_reward_handles_equivalent_latex_and_units() -> None: reward = RewardMathVerify.Config().build() sample = DapoMathSample(prompt="problem", ground_truth=r"336^\circ") - assert asyncio.run(reward(_rollout("work\nAnswer: $336$"), sample)) == 1.0 - assert asyncio.run(reward(_rollout("work\nAnswer: $335$"), sample)) == 0.0 + assert asyncio.run(reward(_rollout(r"work\nAnswer: \boxed{336}"), sample)) == 1.0 + assert asyncio.run(reward(_rollout(r"work\nAnswer: \boxed{335}"), sample)) == 0.0 diff --git a/torchtitan/experiments/rl/tests/test_generator.py b/torchtitan/experiments/rl/tests/test_generator.py index c671363b9e..333272a4a5 100644 --- a/torchtitan/experiments/rl/tests/test_generator.py +++ b/torchtitan/experiments/rl/tests/test_generator.py @@ -32,6 +32,7 @@ from torchtitan.components.checkpointer import CheckpointManager from torchtitan.config import CommConfig, DebugConfig from torchtitan.distributed import utils as dist_utils +from torchtitan.distributed.activation_checkpoint import FullAC from torchtitan.experiments.rl.actors.generator import ( _extract_request_metrics_inputs, _prepare_generation_request_metrics, @@ -315,6 +316,25 @@ def test_trainer_requires_prefix_cache_reset_when_hotswap_off(): ) +def test_qwen36_27b_config_applies_offset_rmsnorm_to_both_actors(): + from torchtitan.experiments.rl.examples.alphabet_sort.config_registry import ( + rl_grpo_qwen3_6_27b_varlen_perf, + ) + + config = rl_grpo_qwen3_6_27b_varlen_perf() + override_import = "torchtitan.overrides.offset_rmsnorm.triton_offset_rmsnorm" + + assert config.hf_assets_path.endswith("Qwen3.6-27B") + assert config.trainer.override.imports == [override_import] + assert config.generator.override.imports == [override_import] + assert config.trainer.parallelism.data_parallel_shard_degree == 2 + assert config.trainer.parallelism.tensor_parallel_degree == 2 + assert config.generator.parallelism.tensor_parallel_degree == 4 + assert config.trainer.optimizer.implementation == "fused_opt_states_bf16" + assert isinstance(config.trainer.ac_config, FullAC.Config) + assert config.generator.cudagraph.enable + + # --- CUDA graph config (VLLMCudagraphConfig.get_vllm_compilation_config) --- diff --git a/torchtitan/experiments/rl/tests/test_muse_glimmer_renderer.py b/torchtitan/experiments/rl/tests/test_muse_glimmer_renderer.py index 9a0410d845..bacdeec3f8 100644 --- a/torchtitan/experiments/rl/tests/test_muse_glimmer_renderer.py +++ b/torchtitan/experiments/rl/tests/test_muse_glimmer_renderer.py @@ -23,6 +23,7 @@ import pytest +from torchtitan.components.tokenizer import HuggingFaceTokenizer from torchtitan.experiments.rl.models.muse_glimmer.renderer import ( EOM_ID, EOT_ID, @@ -227,6 +228,46 @@ def _renderer(tokenizer, **overrides): return MuseGlimmerRenderer(tokenizer, MuseGlimmerRendererConfig(**overrides)) +@pytest.mark.parametrize( + ("kwargs", "error"), + [ + ({"thinking_retention": "everything"}, ValueError), + ({"retain_reasoning_in_history": "false"}, TypeError), + ({"answer_from_reasoning_fallback": 1}, TypeError), + ({"reasoning_strength": 123}, TypeError), + ], +) +def test_config_rejects_invalid_values(kwargs, error): + with pytest.raises(error): + MuseGlimmerRendererConfig(**kwargs) + + +def test_build_snapshots_config(): + path = os.environ.get("MUSE_GLIMMER_TOKENIZER", DEFAULT_TOKENIZER) + if not os.path.isdir(path): + pytest.skip("HuggingFaceTokenizer needs a local tokenizer directory") + config = MuseGlimmerRendererConfig(retain_reasoning_in_history=True) + renderer = config.build(tokenizer=HuggingFaceTokenizer(tokenizer_path=path)) + config.retain_reasoning_in_history = False + assert renderer._config.retain_reasoning_in_history is True + assert renderer.effective_thinking_retention == "all" + + +def test_config_build_matches_hf_tokenizer_path(tokenizer): + path = os.environ.get("MUSE_GLIMMER_TOKENIZER", DEFAULT_TOKENIZER) + if not os.path.isdir(path): + pytest.skip("HuggingFaceTokenizer needs a local tokenizer directory") + titan = MuseGlimmerRendererConfig(reasoning_strength="low").build( + tokenizer=HuggingFaceTokenizer(tokenizer_path=path) + ) + assert isinstance(titan, MuseGlimmerRenderer) + hf = _renderer(tokenizer, reasoning_strength="low") + messages = [{"role": "user", "content": "search for bob"}] + assert titan.render_ids( + messages, tools=TOOLS, add_generation_prompt=True + ) == hf.render_ids(messages, tools=TOOLS, add_generation_prompt=True) + + def _template_kwargs(config: MuseGlimmerRendererConfig) -> dict: """Only forward knobs the caller set, so the template applies its own defaults.""" return { diff --git a/torchtitan/experiments/rl/tests/test_renderer.py b/torchtitan/experiments/rl/tests/test_renderer.py new file mode 100644 index 0000000000..ac99a0caa5 --- /dev/null +++ b/torchtitan/experiments/rl/tests/test_renderer.py @@ -0,0 +1,132 @@ +# 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 json +from dataclasses import dataclass + +import pytest +from renderers import ( + AutoRendererConfig, + DefaultRendererConfig, + OffsetTokenizer, + Qwen3RendererConfig, + Tokenizer, +) + +from torchtitan.components.tokenizer import HuggingFaceTokenizer +from torchtitan.config import Configurable +from torchtitan.experiments.rl.renderer import ( + RendererConfig, + RenderersLibraryConfig, + RendererTokenizerWrapper, +) + +_TOKENIZER_PATH = "tests/assets/tokenizer" + + +# --- RenderersLibraryConfig --- + + +def test_build_renders_with_titan_tokenizer() -> None: + tokenizer = HuggingFaceTokenizer(tokenizer_path=_TOKENIZER_PATH) + renderer = RenderersLibraryConfig( + renderers_config=Qwen3RendererConfig(enable_thinking=False) + ).build(tokenizer=tokenizer) + rendered = renderer.render( + [{"role": "user", "content": "hi"}], add_generation_prompt=True + ) + assert rendered.token_ids[0] == tokenizer.token_to_id("<|im_start|>") + assert renderer.get_stop_token_ids() == [ + tokenizer.token_to_id("<|im_end|>"), + tokenizer.token_to_id("<|endoftext|>"), + ] + assert len(rendered.is_content) == len(rendered.token_ids) + + +@pytest.mark.parametrize( + ("renderers_config", "reason"), + [ + (AutoRendererConfig(), "MODEL_RENDERER_MAP"), + (DefaultRendererConfig(), "special-token variables"), + ], +) +def test_auto_and_default_are_refused(renderers_config, reason: str) -> None: + tokenizer = HuggingFaceTokenizer(tokenizer_path=_TOKENIZER_PATH) + with pytest.raises(ValueError) as error: + RenderersLibraryConfig(renderers_config=renderers_config).build( + tokenizer=tokenizer + ) + assert reason in str(error.value) + assert "Pick the model's renderer" in str(error.value) + + +def test_config_to_dict_is_json() -> None: + # The controller logs `Controller.Config.to_dict()` as the job config. + @dataclass(kw_only=True, slots=True) + class _Holder(Configurable.Config): + renderer: RendererConfig + + holder = _Holder( + renderer=RenderersLibraryConfig( + renderers_config=Qwen3RendererConfig(enable_thinking=False) + ) + ) + value = holder.to_dict() + assert value["renderer"]["renderers_config"]["name"] == "qwen3" + assert value["renderer"]["renderers_config"]["enable_thinking"] is False + json.dumps(value) + + +# --- RendererTokenizerWrapper --- + + +def test_renderer_tokenizer_satisfies_offset_protocol() -> None: + tokenizer = HuggingFaceTokenizer(tokenizer_path=_TOKENIZER_PATH) + renderer_tokenizer = RendererTokenizerWrapper(tokenizer) + assert isinstance(renderer_tokenizer, Tokenizer) + assert isinstance(renderer_tokenizer, OffsetTokenizer) + assert renderer_tokenizer.eos_token_id == tokenizer.eos_id + assert renderer_tokenizer.convert_tokens_to_ids( + "<|im_end|>" + ) == tokenizer.token_to_id("<|im_end|>") + encoding = renderer_tokenizer( + "hi there", add_special_tokens=False, return_offsets_mapping=True + ) + assert encoding["input_ids"] == renderer_tokenizer.encode("hi there") + assert len(encoding["offset_mapping"]) == len(encoding["input_ids"]) + + +def test_encode_never_adds_bos() -> None: + # The debug tokenizer has a BOS token; renderers place special tokens themselves. + tokenizer = HuggingFaceTokenizer(tokenizer_path=_TOKENIZER_PATH) + assert tokenizer.bos_id is not None + assert tokenizer.bos_id not in RendererTokenizerWrapper(tokenizer).encode("hi") + + +def test_render_matches_hf_tokenizer_path() -> None: + transformers = pytest.importorskip("transformers") + from renderers import create_renderer + + messages = [ + {"role": "system", "content": "Sort names."}, + {"role": "user", "content": "Zed, Amy <|im_end|> tricky"}, + {"role": "assistant", "reasoning_content": "think", "content": "Amy, Zed"}, + {"role": "user", "content": "Add Bob."}, + ] + config = Qwen3RendererConfig(enable_thinking=False) + hf = create_renderer( + transformers.AutoTokenizer.from_pretrained(_TOKENIZER_PATH), config + ) + titan = create_renderer( + RendererTokenizerWrapper(HuggingFaceTokenizer(tokenizer_path=_TOKENIZER_PATH)), + config, + ) + expected = hf.render(messages, add_generation_prompt=True) + actual = titan.render(messages, add_generation_prompt=True) + assert actual.token_ids == expected.token_ids + assert actual.is_content == expected.is_content + assert actual.sampled_mask == expected.sampled_mask + assert actual.message_indices == expected.message_indices diff --git a/torchtitan/experiments/rl/tests/test_rollout_worker.py b/torchtitan/experiments/rl/tests/test_rollout_worker.py index ec5fba8c09..d8a9dafc68 100644 --- a/torchtitan/experiments/rl/tests/test_rollout_worker.py +++ b/torchtitan/experiments/rl/tests/test_rollout_worker.py @@ -9,8 +9,13 @@ import asyncio from types import SimpleNamespace +from renderers import Qwen3RendererConfig + +from torchtitan.components.tokenizer import HuggingFaceTokenizer + from torchtitan.experiments.rl.actors.generator import SamplingConfig from torchtitan.experiments.rl.environment.token import TokenEnvOutput +from torchtitan.experiments.rl.renderer import RenderersLibraryConfig from torchtitan.experiments.rl.rollout import RolloutStatus from torchtitan.experiments.rl.rollout.rollouter import RolloutWorker from torchtitan.experiments.rl.rubrics import RubricOutput @@ -121,8 +126,11 @@ async def run() -> None: ) worker = _CustomWorker(worker_config) await worker.setup_async( - renderer_config=_Config("renderer"), - hf_assets_path="hf_assets_path", + tokenizer_config=HuggingFaceTokenizer.Config(), + renderer_config=RenderersLibraryConfig( + renderers_config=Qwen3RendererConfig(enable_thinking=False) + ), + hf_assets_path="tests/assets/tokenizer", ) group = await worker.run_group( generate_fn=generate_fn, @@ -142,7 +150,10 @@ async def run() -> None: assert [rollout.reward for rollout in group.rollouts] == [1.0, 2.0] assert [rollout.advantage for rollout in group.rollouts] == [10.0, 20.0] assert all(env.closed for env in token_env_config.envs) - assert token_env_config.renderers == ["renderer", "renderer"] + assert [type(r).__name__ for r in token_env_config.renderers] == [ + "Qwen3Renderer", + "Qwen3Renderer", + ] assert [call[1]["request_id"] for call in generate_fn.calls] == [ "group=7/rollout=0/turn=0", "group=7/rollout=1/turn=0", diff --git a/torchtitan/experiments/rl/tests/test_rollouter_pool.py b/torchtitan/experiments/rl/tests/test_rollouter_pool.py index a4dad20494..8bd55f5bbb 100644 --- a/torchtitan/experiments/rl/tests/test_rollouter_pool.py +++ b/torchtitan/experiments/rl/tests/test_rollouter_pool.py @@ -96,6 +96,7 @@ async def _setup( host = _ControllerHost(worker_mesh) monkeypatch.setattr(rollouter_module, "this_host", lambda: host) await rollouter.setup_async( + tokenizer_config="tokenizer_config", renderer_config="renderer_config", hf_assets_path="hf_assets_path", ) @@ -110,6 +111,7 @@ async def run() -> None: rollouter = _rollouter_without_datasets() await rollouter.setup_async( + tokenizer_config="tokenizer_config", renderer_config="renderer_config", hf_assets_path="hf_assets_path", ) @@ -125,6 +127,7 @@ async def run() -> None: } assert worker_mesh.actor_mesh.setup_async.calls == [ { + "tokenizer_config": "tokenizer_config", "renderer_config": "renderer_config", "hf_assets_path": "hf_assets_path", } diff --git a/torchtitan/experiments/rl/tests/test_shutdown.py b/torchtitan/experiments/rl/tests/test_shutdown.py index 90d3db1fa2..034729a5e0 100644 --- a/torchtitan/experiments/rl/tests/test_shutdown.py +++ b/torchtitan/experiments/rl/tests/test_shutdown.py @@ -175,11 +175,11 @@ class _StubConfig: rollout_recorder = RolloutSampleRecorder.Config() hf_assets_path = "./tests/assets/tokenizer" # __init__ builds these too; stub them so construction does no real work. + tokenizer = SimpleNamespace( + build=lambda *, tokenizer_path: SimpleNamespace(eos_id=0) + ) renderer = SimpleNamespace( - build=lambda *, tokenizer_path: SimpleNamespace( - get_stop_token_ids=lambda: [], - _tokenizer=SimpleNamespace(eos_token_id=0), - ) + build=lambda *, tokenizer: SimpleNamespace(get_stop_token_ids=lambda: []) ) # __init__ reads generator.sampling (a dataclass, for replace) + generator.debug.seed. generator = SimpleNamespace( diff --git a/torchtitan/experiments/rl/tests/test_token_env.py b/torchtitan/experiments/rl/tests/test_token_env.py new file mode 100644 index 0000000000..bf36db9802 --- /dev/null +++ b/torchtitan/experiments/rl/tests/test_token_env.py @@ -0,0 +1,96 @@ +# 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 asyncio + +from renderers import ParsedResponse + +from torchtitan.experiments.rl.environment.message import ( + MessageEnvInitOutput, + MessageEnvStepOutput, +) +from torchtitan.experiments.rl.environment.token import TokenEnv +from torchtitan.experiments.rl.rollout import RolloutStatus +from torchtitan.experiments.rl.types import Completion + +_TOOLS = [ + { + "type": "function", + "function": { + "name": "search", + "parameters": { + "type": "object", + "properties": {"query": {"type": "string"}}, + }, + }, + } +] + + +class _RecordingRenderer: + """Records the kwargs of each renderer call the env makes.""" + + def __init__(self) -> None: + self.calls: dict[str, dict] = {} + + def render_ids(self, messages, *, tools=None, add_generation_prompt=False): + self.calls["render_ids"] = {"tools": tools} + return [1, 2, 3] + + def parse_response(self, token_ids, *, tools=None): + self.calls["parse_response"] = {"token_ids": token_ids, "tools": tools} + return ParsedResponse(content="answer", reasoning_content=None, tool_calls=[]) + + def bridge_to_next_turn( + self, previous_prompt_ids, previous_completion_ids, new_messages, *, tools=None + ): + self.calls["bridge_to_next_turn"] = {"tools": tools} + return None + + def get_stop_token_ids(self): + return [] + + +class _ToolMessageEnv: + async def init(self) -> MessageEnvInitOutput: + return MessageEnvInitOutput( + init_prompt_messages=[{"role": "user", "content": "find bob"}], + tools=_TOOLS, + ) + + async def step(self, completion_message) -> MessageEnvStepOutput: + return MessageEnvStepOutput( + env_messages=[{"role": "tool", "name": "search", "content": "bob: found"}] + ) + + async def close(self) -> None: + pass + + +def test_env_passes_tools_to_render_parse_and_bridge() -> None: + # Tool schemas are part of the chat template, and XML-style tool parsers need them + # to type the arguments; every renderer call must see the same list. + renderer = _RecordingRenderer() + env = TokenEnv.Config().build(message_env=_ToolMessageEnv(), renderer=renderer) + + async def run(): + await env.init() + return await env.step( + Completion( + min_policy_version=0, + max_policy_version=0, + request_id="r0", + token_ids=[7, 8], + token_logprobs=[-0.1, -0.2], + finish_reason="stop", + ) + ) + + env_output = asyncio.run(run()) + assert env_output.status == RolloutStatus.ONGOING + assert renderer.calls["render_ids"]["tools"] == _TOOLS + assert renderer.calls["parse_response"] == {"token_ids": [7, 8], "tools": _TOOLS} + assert renderer.calls["bridge_to_next_turn"]["tools"] == _TOOLS diff --git a/torchtitan/experiments/torchft/trainer.py b/torchtitan/experiments/torchft/trainer.py index 11875e14a5..666b5fd807 100644 --- a/torchtitan/experiments/torchft/trainer.py +++ b/torchtitan/experiments/torchft/trainer.py @@ -56,6 +56,7 @@ def __init__(self, config: Config): # init distributed and build meshes (FT override handles ft_manager creation) self.parallel_dims = parallel_dims = self.init_distributed() + dist_utils.set_spmd_backend(config.parallelism.spmd_backend) # Logging needs to happen after distributed initialized config.maybe_log() @@ -303,6 +304,10 @@ def __init__(self, config: Config): self.train_context = dist_utils.get_spmd_context( parallel_dims=parallel_dims, + spmd_typechecking=( + config.parallelism.spmd_backend == "spmd_types" + and config.debug.spmd_typechecking + ), ) self.fwd_bwd_fn = self._forward_backward_body if not config.training.disable_cuda_graphs: diff --git a/torchtitan/models/common/config_utils.py b/torchtitan/models/common/config_utils.py index 0bab03b9bc..5290c7de5a 100644 --- a/torchtitan/models/common/config_utils.py +++ b/torchtitan/models/common/config_utils.py @@ -360,8 +360,8 @@ def make_token_dispatcher_config( num_experts: int, top_k: int, comm_backend: str, + hidden_dim: int, non_blocking_capacity_factor: float | None = None, - hidden_dim: int | None = None, num_max_tokens_per_rank: int | None = None, cudagraphable: bool = False, ) -> LocalTokenDispatcher.Config: @@ -409,6 +409,7 @@ def make_token_dispatcher_config( return MinimalAsyncEPTokenDispatcher.Config( num_experts=num_experts, top_k=top_k, + hidden_dim=hidden_dim, num_max_tokens_per_rank=num_max_tokens_per_rank, ) elif comm_backend == "standard": diff --git a/torchtitan/models/common/moe.py b/torchtitan/models/common/moe.py index f9c1fa13ca..f5b207b439 100644 --- a/torchtitan/models/common/moe.py +++ b/torchtitan/models/common/moe.py @@ -29,7 +29,11 @@ # e = num local experts (E / EP, used in token dispatcher for # per-local-expert token counts after EP dispatch /_permute), # K = top-k, N = routed tokens (T*K), -# R = routed tokens assigned to local experts +# R = routed tokens assigned to local experts, +# O = expert output features, I = expert input features +# (roles, not model dims: the _grouped_mm seam takes the expert +# weight in its stored (E, O, I) orientation, which is (E, F, D) +# for the up/gate projections and (E, D, F) for the down one) class GroupedExperts(Module): diff --git a/torchtitan/models/common/token_dispatcher.py b/torchtitan/models/common/token_dispatcher.py index b5edd242f0..459883b431 100644 --- a/torchtitan/models/common/token_dispatcher.py +++ b/torchtitan/models/common/token_dispatcher.py @@ -1232,7 +1232,6 @@ def update_ep_token_dispatcher_config(model_config: Any, config: Any) -> None: "parallelism (expert_parallel_degree > 1)." ) - token_dispatcher_cfg.hidden_dim = model_config.dim configured_capacity = token_dispatcher_cfg.num_max_tokens_per_rank if configured_capacity is not None and configured_capacity <= 0: raise ValueError( diff --git a/torchtitan/models/deepseek_v3/__init__.py b/torchtitan/models/deepseek_v3/__init__.py index edb6552417..e480ff02d4 100644 --- a/torchtitan/models/deepseek_v3/__init__.py +++ b/torchtitan/models/deepseek_v3/__init__.py @@ -118,7 +118,7 @@ def make_mla_attention_config( wq_b = None # q_norm is unused when q_lora_rank == 0 (never built), but the field is # required on Attention.Config so we supply a placeholder. - q_norm = RMSNorm.Config(normalized_shape=1, param_init=norm_init) + q_norm = RMSNorm.Config(normalized_shape=1, eps=1e-6, param_init=norm_init) else: wq = None wq_a = Linear.Config( @@ -131,7 +131,9 @@ def make_mla_attention_config( out_features=n_heads * qk_head_dim, param_init=linear_init, ) - q_norm = RMSNorm.Config(normalized_shape=q_lora_rank, param_init=norm_init) + q_norm = RMSNorm.Config( + normalized_shape=q_lora_rank, eps=1e-6, param_init=norm_init + ) return Attention.Config( dim=dim, @@ -151,7 +153,9 @@ def make_mla_attention_config( out_features=kv_lora_rank + qk_rope_head_dim, param_init=linear_init, ), - kv_norm=RMSNorm.Config(normalized_shape=kv_lora_rank, param_init=norm_init), + kv_norm=RMSNorm.Config( + normalized_shape=kv_lora_rank, eps=1e-6, param_init=norm_init + ), wkv_b=Linear.Config( in_features=kv_lora_rank, out_features=n_heads * (qk_nope_head_dim + v_head_dim), @@ -271,9 +275,11 @@ def build_mla_moe_layers( DeepSeekV3TransformerBlock.Config( attention=attn_cfg, attention_norm=RMSNorm.Config( - normalized_shape=dim, param_init=norm_init + normalized_shape=dim, eps=1e-6, param_init=norm_init + ), + ffn_norm=RMSNorm.Config( + normalized_shape=dim, eps=1e-6, param_init=norm_init ), - ffn_norm=RMSNorm.Config(normalized_shape=dim, param_init=norm_init), feed_forward=ffn_cfg, moe=moe_cfg, ) @@ -307,14 +313,14 @@ def _build_mtp_layers( moe=copy.deepcopy(inner_cfg.moe), attention_norm=copy.deepcopy(inner_cfg.attention_norm), ffn_norm=copy.deepcopy(inner_cfg.ffn_norm), - enorm=RMSNorm.Config(normalized_shape=dim), - hnorm=RMSNorm.Config(normalized_shape=dim), + enorm=RMSNorm.Config(normalized_shape=dim, eps=1e-6), + hnorm=RMSNorm.Config(normalized_shape=dim, eps=1e-6), eh_proj=Linear.Config( in_features=dim * 2, out_features=dim, bias=False, ), - mtp_norm=RMSNorm.Config(normalized_shape=dim), + mtp_norm=RMSNorm.Config(normalized_shape=dim, eps=1e-6), ) ) return mtp_layers @@ -376,7 +382,7 @@ def _debugmodel( tok_embeddings=Embedding.Config( num_embeddings=vocab_size, embedding_dim=dim, param_init=_EMBEDDING_INIT ), - norm=RMSNorm.Config(normalized_shape=dim, param_init=_NORM_INIT), + norm=RMSNorm.Config(normalized_shape=dim, eps=1e-6, param_init=_NORM_INIT), lm_head=Linear.Config( in_features=dim, out_features=vocab_size, @@ -448,7 +454,7 @@ def _16b( tok_embeddings=Embedding.Config( num_embeddings=vocab_size, embedding_dim=dim, param_init=_EMBEDDING_INIT ), - norm=RMSNorm.Config(normalized_shape=dim, param_init=_NORM_INIT), + norm=RMSNorm.Config(normalized_shape=dim, eps=1e-6, param_init=_NORM_INIT), lm_head=Linear.Config( in_features=dim, out_features=vocab_size, @@ -524,7 +530,7 @@ def _236b( tok_embeddings=Embedding.Config( num_embeddings=vocab_size, embedding_dim=dim, param_init=_EMBEDDING_INIT ), - norm=RMSNorm.Config(normalized_shape=dim, param_init=_NORM_INIT), + norm=RMSNorm.Config(normalized_shape=dim, eps=1e-6, param_init=_NORM_INIT), lm_head=Linear.Config( in_features=dim, out_features=vocab_size, @@ -601,7 +607,7 @@ def _671b( tok_embeddings=Embedding.Config( num_embeddings=vocab_size, embedding_dim=dim, param_init=_EMBEDDING_INIT ), - norm=RMSNorm.Config(normalized_shape=dim, param_init=_NORM_INIT), + norm=RMSNorm.Config(normalized_shape=dim, eps=1e-6, param_init=_NORM_INIT), lm_head=Linear.Config( in_features=dim, out_features=vocab_size, diff --git a/torchtitan/models/deepseek_v3/config_registry.py b/torchtitan/models/deepseek_v3/config_registry.py index 0cf1ab695e..9953b96b48 100644 --- a/torchtitan/models/deepseek_v3/config_registry.py +++ b/torchtitan/models/deepseek_v3/config_registry.py @@ -25,6 +25,30 @@ from . import model_registry +def deepseek_v3_mxfp8_linear_converter_config( + *, model_compile_enabled: bool +) -> MXFP8LinearConverter.Config: + """Build the dense MXFP8 policy shared by eager and GraphTrainer configs. + + The KV up projection and FFN down projections have single-consumer inputs + that are not saved elsewhere for backward, so their columnwise MXFP8 + representations replace BF16 storage. Shared-input and attention output + projections use the conservative BF16 save format. This selection is based + on activation ownership, not the activation-checkpointing policy. + Checkpointing changes when the selected representation is recreated and how + long it remains live. + """ + return MXFP8LinearConverter.Config( + model_compile_enabled=model_compile_enabled, + fqns=["attention", "shared_experts", "feed_forward"], + linears_saving_inputs_for_backward_in_mxfp8=[ + "attention.wkv_b", + "feed_forward.w2", + "shared_experts.w2", + ], + ) + + def enable_fused_swiglu(config: Trainer.Config) -> None: # Activate the stock dense-FFN and MoE grouped-expert overrides. The separate # dist-GEMM FFN override is not needed by these configs. @@ -97,9 +121,8 @@ def deepseek_v3_debugmodel_mxfp8(seq_len: int | None = None) -> Trainer.Config: "debugmodel", seq_len=seq_len, converters=[ - MXFP8LinearConverter.Config( + deepseek_v3_mxfp8_linear_converter_config( model_compile_enabled=model_compile_enabled, - fqns=["attention", "shared_experts", "feed_forward"], ), MXFP8GroupedExpertsConverter.Config( model_compile_enabled=model_compile_enabled, diff --git a/torchtitan/models/deepseek_v4/model.py b/torchtitan/models/deepseek_v4/model.py index fa38e3af59..ec8c94522d 100644 --- a/torchtitan/models/deepseek_v4/model.py +++ b/torchtitan/models/deepseek_v4/model.py @@ -5,13 +5,18 @@ # LICENSE file in the root directory of this source tree. from dataclasses import dataclass -from typing import TYPE_CHECKING +from typing import cast, TYPE_CHECKING import torch +from torch import nn from torchtitan.models.common.attention import AttentionMasksType from torchtitan.models.common.decoder import Decoder, TransformerBlock from torchtitan.models.deepseek_v3.mtp import roll_mtp_sequence +from torchtitan.models.utils import ( + get_nparams_and_active_nparams, + quadratic_attention_flops_per_token, +) from .mhc import HcHead, HcPost, HcPre @@ -144,20 +149,55 @@ def update_from_config(self, *, config, **kwargs): enable_ep=parallelism.expert_parallel_degree > 1, ) - def get_nparams_and_flops(self, model, seq_len): - total_params = sum(p.numel() for p in model.parameters() if p.requires_grad) - non_embed_params = sum( - p.numel() - for n, p in model.named_parameters() - if p.requires_grad and "tok_embeddings" not in n and "lm_head" not in n + def get_nparams_and_flops( + self, model: nn.Module, seq_len: int + ) -> tuple[int, int]: + """Estimate DeepSeek V4 training FLOPs from the final model config.""" + deepseek_v4_model = cast(DeepSeekV4Model, model) + nparams, active_nparams = get_nparams_and_active_nparams(deepseek_v4_model) + + attention_op_flops = 0 + for layers in (self.layers, self.mtp_layers or ()): + for layer in layers: + attention = layer.attention + inner_attention = attention.inner_attention + attention_op_flops += quadratic_attention_flops_per_token( + num_heads=attention.n_heads, + qk_head_dim=attention.head_dim, + v_head_dim=attention.head_dim, + seq_len=seq_len, + sliding_window_size=inner_attention.window_size, + ) + + if attention.compress_ratio > 1: + compressed_seq_len = seq_len // attention.compress_ratio + if attention.compress_ratio == 4: + attention_op_flops += ( + 6 + * attention.index_n_heads + * attention.index_head_dim + * compressed_seq_len + ) + compressed_seq_len = min( + compressed_seq_len, inner_attention.index_topk + ) + attention_op_flops += quadratic_attention_flops_per_token( + num_heads=attention.n_heads, + qk_head_dim=attention.head_dim, + v_head_dim=attention.head_dim, + seq_len=compressed_seq_len, + ) + + active_nparams += len(deepseek_v4_model.mtp_layers) * sum( + param.numel() for param in deepseek_v4_model.lm_head.parameters() ) - n_layers = self.n_layers + self.n_mtp_layers - head_dim = self.layers[0].attention.head_dim - n_heads = self.layers[0].attention.n_heads - flops_per_token = ( - 6 * non_embed_params + 12 * n_layers * n_heads * head_dim * seq_len + active_nparams += (self.hc_mult - 1) * sum( + param.numel() + for mtp_layer in deepseek_v4_model.mtp_layers + for param in cast("MTPBlock", mtp_layer).h_proj.parameters() ) - return total_params, int(flops_per_token) + + return nparams, 6 * active_nparams + attention_op_flops def __init__(self, config: Config): super().__init__(config) diff --git a/torchtitan/models/flux/README.md b/torchtitan/models/flux/README.md index 822bbf150a..e481051237 100644 --- a/torchtitan/models/flux/README.md +++ b/torchtitan/models/flux/README.md @@ -110,4 +110,3 @@ The `fqns` parameter specifies which fully qualified module names to quantize. T ## TODO - [ ] More parallelism support (Tensor Parallelism, Pipeline Parallelism, etc) -- [ ] Implement the num_flops_per_token calculation in get_nparams_and_flops() function diff --git a/torchtitan/models/gpt_oss/README.md b/torchtitan/models/gpt_oss/README.md index 68bbc77310..1c92354929 100644 --- a/torchtitan/models/gpt_oss/README.md +++ b/torchtitan/models/gpt_oss/README.md @@ -6,9 +6,15 @@ MODULE=gpt_oss CONFIG=gpt_oss_debugmodel ./run_train.sh ``` ## Supported Features -- FSDP/HSDP, TP, EP +- FSDP/HSDP, TP, EP, CP, PP - Grouped matrix multiplication for efficient computation +CI already runs CP and PP on the debug model: +- `gpt_oss_pp+fsdp+cp+ep+sacop` (`gpt_oss_debugmodel_flex_fsdp2_cp2_pp2_ep4_sac`) +- `gpt_oss_pp+fsdp+ep+sacop` -## TODO -1. More parallelism support: CP, PP +Those jobs use Interleaved1F1B. FlexAttention zero-bubble / split-backward PP +tests are disabled in `tests/integration_tests/features.py` because +FlexAttention `BlockMask` is not a Tensor (`stage_backward_input` calls +`requires_grad` on every stage input). Full-backward schedules +(1F1B / GPipe / Interleaved1F1B) are unaffected. diff --git a/torchtitan/models/kimi_k3/__init__.py b/torchtitan/models/kimi_k3/__init__.py index 9543063621..9aa824cef2 100644 --- a/torchtitan/models/kimi_k3/__init__.py +++ b/torchtitan/models/kimi_k3/__init__.py @@ -12,10 +12,12 @@ from torchtitan.components.optimizer import register_moe_load_balancing_hook from torchtitan.models.common import Conv1d, Embedding, Linear -from torchtitan.models.common.config_utils import get_attention_config +from torchtitan.models.common.config_utils import ( + get_attention_config, + make_token_dispatcher_config, +) from torchtitan.models.common.moe import RoutedExperts, TokenChoiceTopKRouter from torchtitan.models.common.nn_modules import GELU, RMSNorm -from torchtitan.models.common.token_dispatcher import LocalTokenDispatcher from torchtitan.models.common.vision_encoder import ( VisionAttention, VisionMLP, @@ -222,6 +224,7 @@ def _latent_moe_config( num_experts: int, top_k: int, num_shared_experts: int, + moe_comm_backend: str, ) -> KimiLatentMoE.Config: return KimiLatentMoE.Config( num_experts=num_experts, @@ -247,9 +250,16 @@ def _latent_moe_config( "w3_EFD": partial(nn.init.trunc_normal_, std=0.02), }, ), - token_dispatcher=LocalTokenDispatcher.Config( + # core's dispatcher factory: standard / deepep / hybridep / + # minimal_async_ep per spec, as deepseek_v3; falls back to local + # dispatch when the ep mesh is None. + token_dispatcher=make_token_dispatcher_config( num_experts=num_experts, top_k=top_k, + comm_backend=moe_comm_backend, + # The routed experts consume the LATENT stream, so the + # dispatcher buffers size by latent_dim, not model dim. + hidden_dim=latent_dim, ), ), routed_norm=_norm(latent_dim), @@ -369,6 +379,7 @@ def _kimi_k3_config( num_shared_experts: int, vision_encoder: KimiK3VisionEncoder.Config, attn_backend: str, + moe_comm_backend: str = "standard", ) -> KimiK3Model.Config: """Assemble a Kimi K3 config from the released topology's free parameters. @@ -422,6 +433,7 @@ def _kimi_k3_config( num_experts=num_experts, top_k=top_k, num_shared_experts=num_shared_experts, + moe_comm_backend=moe_comm_backend, ) ), attention_norm=_norm(dim), @@ -454,10 +466,11 @@ def _kimi_k3_config( ) -def _debugmodel(attn_backend: str) -> KimiK3Model.Config: +def _debugmodel(attn_backend: str, moe_comm_backend: str) -> KimiK3Model.Config: dim = 1024 return _kimi_k3_config( dim=dim, + moe_comm_backend=moe_comm_backend, vocab_size=163840, num_layers=24, full_attention_layers={3, 7, 11, 15, 19, 23}, @@ -490,10 +503,11 @@ def _debugmodel(attn_backend: str) -> KimiK3Model.Config: ) -def _kimi_k3(attn_backend: str) -> KimiK3Model.Config: +def _kimi_k3(attn_backend: str, moe_comm_backend: str) -> KimiK3Model.Config: dim = 7168 return _kimi_k3_config( dim=dim, + moe_comm_backend=moe_comm_backend, vocab_size=163840, num_layers=93, full_attention_layers=set(range(3, 92, 4)) | {92}, @@ -536,6 +550,7 @@ def model_registry( flavor: str, attn_backend: str = "flex", converters: list[ModelConfigConverter.Config] | None = None, + moe_comm_backend: str = "standard", *, seq_len: int | None = None, ) -> ModelSpec: @@ -548,7 +563,7 @@ def model_registry( f"Requested seq_len {context_len} exceeds max context length " f"{max_context_len} for flavor {flavor}" ) - config = get_config(attn_backend=attn_backend) + config = get_config(attn_backend=attn_backend, moe_comm_backend=moe_comm_backend) if converters is not None: validate_converter_order(converters) for converter in converters: diff --git a/torchtitan/models/kimi_k3/model.py b/torchtitan/models/kimi_k3/model.py index 8dce8c44df..2280576653 100644 --- a/torchtitan/models/kimi_k3/model.py +++ b/torchtitan/models/kimi_k3/model.py @@ -23,6 +23,7 @@ scatter_vision_embeds, ) from torchtitan.models.common.nn_modules import RMSNorm +from torchtitan.models.kimi_k3.sharding import set_kimi_k3_sharding_config from torchtitan.models.utils import ( delta_rule_flops_per_token, get_nparams_and_active_nparams, @@ -277,6 +278,9 @@ def update_from_config(self, *, config, **kwargs) -> None: # and KDA recurrent states at document boundaries. if isinstance(dataset, MMSamplePackingConfig): raise ValueError("Kimi K3 does not yet support sample packing.") + set_kimi_k3_sharding_config( + self, enable_ep=config.parallelism.expert_parallel_degree > 1 + ) Decoder.Config.update_from_config(self, config=config, **kwargs) def get_nparams_and_flops( diff --git a/torchtitan/models/kimi_k3/parallelize.py b/torchtitan/models/kimi_k3/parallelize.py index 8a7d604a02..2565b07bde 100644 --- a/torchtitan/models/kimi_k3/parallelize.py +++ b/torchtitan/models/kimi_k3/parallelize.py @@ -39,7 +39,6 @@ def parallelize_kimi_k3( ("tensor parallel", parallel_dims.tp_enabled), ("pipeline parallel", parallel_dims.pp_enabled), ("context parallel", parallel_dims.cp_enabled), - ("expert parallel", parallel_dims.ep_enabled), ) if enabled ] @@ -60,8 +59,23 @@ def parallelize_kimi_k3( ["dp_replicate", "fsdp"] if parallel_dims.dp_replicate_enabled else ["fsdp"] ) dp_mesh = parallel_dims.get_mesh(dp_mesh_names) + # The routed experts shard on their own data-parallel mesh, which excludes + # the expert axis; the same shape deepseek_v3 resolves. + edp_mesh = None + if parallel_dims.ep_enabled: + edp_mesh = parallel_dims.get_optional_mesh( + ["dp_replicate", "efsdp"] + if parallel_dims.dp_replicate_enabled + else ["efsdp"] + ) assert isinstance(model, KimiK3Model) + if parallel_dims.ep_enabled: + # model_registry's moe_comm_backend picks the dispatcher: standard + # (default), deepep and minimal_async_ep run on this model; hybridep + # needs GB200-class hardware. + model.parallelize(parallel_dims) + if ac_config is not None: ac_policy = ac_config.build(dump_folder=dump_folder) ac_policy.apply(model) @@ -90,7 +104,8 @@ def parallelize_kimi_k3( pp_enabled=False, cpu_offload=training.enable_cpu_offload, reshard_after_forward_policy=parallelism.fsdp_reshard_after_forward, - ep_degree=1, + ep_degree=parallel_dims.ep, + edp_mesh=edp_mesh, enable_symm_mem=parallelism.enable_fsdp_symm_mem, ) diff --git a/torchtitan/models/kimi_k3/sharding.py b/torchtitan/models/kimi_k3/sharding.py new file mode 100644 index 0000000000..fcafe72090 --- /dev/null +++ b/torchtitan/models/kimi_k3/sharding.py @@ -0,0 +1,47 @@ +# 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. + +"""Sharding configs for Kimi K3. Same pattern as ``qwen3_5/sharding.py``. + +Declarations only: functions here set ``ShardingConfig`` on sub-configs of an +already-built config tree, and ``model.parallelize()`` applies them through the +Module protocol. Nothing here touches a mesh or a device. +""" + +from typing import TYPE_CHECKING + +import spmd_types as spmd + +from torchtitan.models.common.moe_sharding import set_moe_sharding_config + +if TYPE_CHECKING: + from torchtitan.models.kimi_k3.model import KimiK3Model + + +def set_kimi_k3_sharding_config( + config: "KimiK3Model.Config", *, enable_ep: bool, enable_sp: bool = False +) -> None: + """Declare the sharding expert parallel acts on. + + The routed experts shard on the expert axis; ``set_moe_sharding_config`` + declares that layout, and its input boundary lifts the plain incoming + activations itself, so no decoder-level declaration is needed. + """ + for layer in config.layers: + if layer.moe is not None: + set_moe_sharding_config( + layer.moe, + enable_ep=enable_ep, + # TODO: flip to True from the caller once the + # tensor-parallel PR lands; with EP alone the internals run + # without sequence parallel. + enable_sp=enable_sp, + expert_param_layout={ + "w1_EFD": spmd.S(1), + "w2_EDF": spmd.S(2), + "w3_EFD": spmd.S(1), + }, + ) diff --git a/torchtitan/models/llama3/README.md b/torchtitan/models/llama3/README.md new file mode 100644 index 0000000000..28ce974b25 --- /dev/null +++ b/torchtitan/models/llama3/README.md @@ -0,0 +1,53 @@ +# Llama 3 + +Llama 3 is the reference decoder model in torchtitan. Training recipes cover +Llama 3.1 8B, 70B, and 405B, plus a small debug model used by CI. + +## Download the tokenizer + +Follow the access instructions on the official +[meta-llama](https://huggingface.co/meta-llama/Llama-3.1-8B) repository, then: + +```bash +python scripts/download_hf_assets.py --repo_id meta-llama/Llama-3.1-8B --assets tokenizer +``` + +The 8B, 70B, and 405B recipes expect tokenizer assets under +`./assets/hf/Llama-3.1-8B` (and the matching 70B / 405B paths). + +## Training + +```bash +# Debug model (used by integration tests) +MODULE=llama3 CONFIG=llama3_debugmodel ./run_train.sh + +# Llama 3.1 8B +MODULE=llama3 CONFIG=llama3_8b ./run_train.sh +``` + +Other recipes include `llama3_70b` and `llama3_405b`. See +[`config_registry.py`](./config_registry.py). + +## Supported Parallelisms + +Coverage below matches `parallelize.py` and the Llama 3 jobs in +`tests/integration_tests/features.py` (plus Float8 jobs in +`tests/integration_tests/h100.py`). + +| Feature | Notes | +|---------|-------| +| FSDP / HSDP | Default data-parallel path | +| Tensor Parallel (TP) | Including sequence parallel; async TP is exercised on H100 | +| Context Parallel (CP) | Composes with FSDP, HSDP, DDP, and TP | +| Pipeline Parallel (PP) | 1F1B, Interleaved1F1B, and GPipe. Zero-bubble / split-backward PP tests are disabled in `tests/integration_tests/features.py` because FlexAttention `BlockMask` is not a Tensor | +| DDP | Including DDP+CP | +| Activation checkpointing | Selective and full | +| `torch.compile` | 1D and multi-dimensional jobs | +| Float8 | H100 integration tests; `llama3_debugmodel_float8` | +| MXFP8 | Recipe `llama3_8b_mxfp8` exists; not in the default GPU feature suite | + +## Numerical checks + +Llama 3 is the baseline used to validate distributed training techniques. See +`tests/integration_tests` and [docs/converging.md](/docs/converging.md) rather +than treating any single published KL or MFU number as a parity claim. diff --git a/torchtitan/models/llama3/config_registry.py b/torchtitan/models/llama3/config_registry.py index 13faf71a1d..deaa14c33c 100644 --- a/torchtitan/models/llama3/config_registry.py +++ b/torchtitan/models/llama3/config_registry.py @@ -35,6 +35,28 @@ from .model import Llama3Model +def llama3_mxfp8_linear_converter_config( + *, model_compile_enabled: bool +) -> MXFP8LinearConverter.Config: + """Build the MXFP8 policy shared by eager and GraphTrainer configs. + + The fused QKV and FFN down projections have single-consumer inputs that are + not saved elsewhere for backward, so their columnwise MXFP8 representations + replace BF16 storage. Other projections retain the conservative BF16 save + format because their inputs are shared or retained elsewhere. This selection + is based on activation ownership, not the activation-checkpointing policy. + Checkpointing changes when the selected representation is recreated and how + long it remains live. + """ + return MXFP8LinearConverter.Config( + model_compile_enabled=model_compile_enabled, + linears_saving_inputs_for_backward_in_mxfp8=[ + "attention.qkv_linear.wqkv", + "feed_forward.w2", + ], + ) + + def llama3_debugmodel(seq_len: int | None = None) -> Trainer.Config: model_spec = model_registry("debugmodel", seq_len=seq_len) packed = ConcatThenSplitPackingConfig(dataset=DATASETS["c4_test"]) @@ -122,6 +144,19 @@ def llama3_debugmodel_float8(seq_len: int | None = None) -> Trainer.Config: return config +def llama3_debugmodel_mxfp8(seq_len: int | None = None) -> Trainer.Config: + config = llama3_debugmodel(seq_len=seq_len) + config.compile = CompileConfig(enable=True, components=["model"]) + config.model_spec = model_registry( + "debugmodel", + seq_len=seq_len, + converters=[ + llama3_mxfp8_linear_converter_config(model_compile_enabled=True), + ], + ) + return config + + def llama3_debugmodel_nvfp4(seq_len: int | None = None) -> Trainer.Config: config = llama3_debugmodel(seq_len=seq_len) config.parallelism.spmd_backend = "spmd_types" @@ -268,7 +303,7 @@ def llama3_8b_mxfp8(seq_len: int | None = None) -> Trainer.Config: "8B", seq_len=seq_len, converters=[ - MXFP8LinearConverter.Config(model_compile_enabled=True), + llama3_mxfp8_linear_converter_config(model_compile_enabled=True), ], ) return config diff --git a/torchtitan/models/qwen3/README.md b/torchtitan/models/qwen3/README.md index 243dee688a..e0b8c9ed57 100644 --- a/torchtitan/models/qwen3/README.md +++ b/torchtitan/models/qwen3/README.md @@ -1,6 +1,3 @@ -**The Qwen3 model is still under development.** - - ## Available features #### Dense Model - Qwen3 dense model: @@ -9,10 +6,15 @@ - Qwen3 MoE model: - Supports FSDP/HSDP, TP, CP, DDP, EP. - Supports AC, torch.compile. - - MoE models use Token Choice routing, which is using auxiluary-loss-free load balancing algorithm. + - MoE models use Token Choice routing. Load-balancing follows the original + Qwen3 training recipe (auxiliary loss). +Dense and MoE debug models are covered by `tests/integration_tests/models.py` +(FSDP+TP+CP, and FSDP+TP+CP+EP for MoE). -Other model sizes are added to the configs, but config_registry entries need to be added and tested. +Model architectures exist for 4B, 8B, and 235B-A22B, but those sizes do not +yet have pretrain `config_registry` recipes. 8B is available as SFT via +`sft_qwen3_8b_math`. ## Download Qwen3 tokenizer ```python scripts/download_hf_assets.py --repo_id --assets tokenizer``` @@ -20,7 +22,8 @@ Other model sizes are added to the configs, but config_registry entries need to eg, for Qwen3 0.6B model, the HF repo name is `Qwen/Qwen3-0.6B`. For 1.7B model, the HF repo name is `Qwen/Qwen3-1.7B`. -## To be added -- Testing - - Learning rate verifying: verify learning rate and schedule with real training jobs (eg, 3k stps), or find official references. - - The model should be tested against established performance benchmarks +## Remaining work +- Add `config_registry` recipes for 4B, 8B pretrain, and 235B-A22B. +- Verify learning rate and schedule on longer training jobs, or cite official + references. +- Compare against established performance benchmarks. diff --git a/torchtitan/overrides/README.md b/torchtitan/overrides/README.md index e1b19891af..1496bf1967 100644 --- a/torchtitan/overrides/README.md +++ b/torchtitan/overrides/README.md @@ -506,6 +506,11 @@ for the full recipe. recipe from "Custom kernels and `torch.compile`". `helion` is an optional dependency, so the module imports without it and falls back to the PyTorch RoPE when it (or CUDA) is unavailable; it is checkpoint-compatible with stock. +- `torchtitan/overrides/offset_rmsnorm.py` — replaces Qwen3.5 + `OffsetRMSNorm` with fused Triton forward and backward kernels while preserving + the stock zero-centered weight and checkpoint layout. Activate it with + `--override.imports torchtitan.overrides.offset_rmsnorm.triton_offset_rmsnorm`. + The `TritonRoPE` snippets above are illustrative — no `triton_rope.py` is shipped — but RoPE is a fully valid override target (`helion_rope.py` is a real one): each attention module owns a `rope` submodule (`RoPE.Config`), so a custom diff --git a/torchtitan/overrides/offset_rmsnorm.py b/torchtitan/overrides/offset_rmsnorm.py new file mode 100644 index 0000000000..218e373390 --- /dev/null +++ b/torchtitan/overrides/offset_rmsnorm.py @@ -0,0 +1,380 @@ +# 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. + +"""Fused Triton OffsetRMSNorm override for Qwen3.5. + +The stock module computes ``(1 + weight) * rmsnorm(input)`` with eager +PyTorch operations. This override keeps that parameterization and performs the +normalization and offset scaling in one Triton kernel. The backward uses Triton +kernels for both the input and weight gradients. +""" + +from __future__ import annotations + +from dataclasses import dataclass, replace + +import torch +import triton +import triton.language as tl +from spmd_types import SpmdType +from torch.distributed.tensor import DTensor + +from torchtitan.config import derive, override +from torchtitan.models.qwen3_5.model import OffsetRMSNorm +from torchtitan.protocols.sharding import LocalMapConfig, resolve_placements + + +__all__ = [ + "TritonOffsetRMSNorm", + "triton_offset_rms_norm", + "triton_offset_rmsnorm", +] + + +_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16, torch.float32) +_MAX_BLOCK_SIZE = 65536 +_DW_BLOCK_M = 32 +_DW_BLOCK_N = 64 + + +def _num_warps(block_size: int) -> int: + if block_size >= 8192: + return 16 + if block_size >= 2048: + return 8 + return 4 + + +@triton.jit +def _offset_rms_norm_forward_kernel( + input_ptr, + weight_ptr, + output_ptr, + inverse_rms_ptr, + num_cols: tl.constexpr, + eps: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +) -> None: + row_idx = tl.program_id(0) + col_offsets = tl.arange(0, BLOCK_SIZE) + mask = col_offsets < num_cols + + input_row = input_ptr + row_idx * num_cols + output_row = output_ptr + row_idx * num_cols + input_fp32 = tl.load(input_row + col_offsets, mask=mask, other=0.0).to(tl.float32) + variance = tl.sum(input_fp32 * input_fp32, axis=0) / num_cols + inverse_rms = tl.rsqrt(variance + eps) + weight_fp32 = tl.load(weight_ptr + col_offsets, mask=mask, other=0.0).to(tl.float32) + output_fp32 = input_fp32 * inverse_rms * (1.0 + weight_fp32) + + tl.store(output_row + col_offsets, output_fp32, mask=mask) + tl.store(inverse_rms_ptr + row_idx, inverse_rms) + + +@triton.jit +def _offset_rms_norm_input_grad_kernel( + grad_output_ptr, + input_ptr, + weight_ptr, + inverse_rms_ptr, + grad_input_ptr, + num_cols: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +) -> None: + row_idx = tl.program_id(0) + col_offsets = tl.arange(0, BLOCK_SIZE) + mask = col_offsets < num_cols + row_offsets = row_idx * num_cols + col_offsets + + grad_output_fp32 = tl.load(grad_output_ptr + row_offsets, mask=mask, other=0.0).to( + tl.float32 + ) + input_fp32 = tl.load(input_ptr + row_offsets, mask=mask, other=0.0).to(tl.float32) + weight_fp32 = tl.load(weight_ptr + col_offsets, mask=mask, other=0.0).to(tl.float32) + inverse_rms = tl.load(inverse_rms_ptr + row_idx) + + scaled_grad = grad_output_fp32 * (1.0 + weight_fp32) + projection = tl.sum(scaled_grad * input_fp32, axis=0) + grad_input_fp32 = inverse_rms * scaled_grad + grad_input_fp32 -= ( + input_fp32 * inverse_rms * inverse_rms * inverse_rms * projection / num_cols + ) + tl.store(grad_input_ptr + row_offsets, grad_input_fp32, mask=mask) + + +@triton.jit +def _offset_rms_norm_weight_grad_partial_kernel( + grad_output_ptr, + input_ptr, + inverse_rms_ptr, + partial_grad_weight_ptr, + num_rows, + num_cols: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +) -> None: + col_block_idx = tl.program_id(0) + row_block_idx = tl.program_id(1) + row_offsets = row_block_idx * BLOCK_M + tl.arange(0, BLOCK_M) + col_offsets = col_block_idx * BLOCK_N + tl.arange(0, BLOCK_N) + offsets = row_offsets[:, None] * num_cols + col_offsets[None, :] + mask = (row_offsets[:, None] < num_rows) & (col_offsets[None, :] < num_cols) + + grad_output_fp32 = tl.load(grad_output_ptr + offsets, mask=mask, other=0.0).to( + tl.float32 + ) + input_fp32 = tl.load(input_ptr + offsets, mask=mask, other=0.0).to(tl.float32) + inverse_rms = tl.load( + inverse_rms_ptr + row_offsets, + mask=row_offsets < num_rows, + other=0.0, + ) + partial = tl.sum( + grad_output_fp32 * input_fp32 * inverse_rms[:, None], + axis=0, + ) + partial_offsets = row_block_idx * num_cols + col_offsets + tl.store( + partial_grad_weight_ptr + partial_offsets, + partial, + mask=col_offsets < num_cols, + ) + + +@triton.jit +def _offset_rms_norm_weight_grad_reduce_kernel( + partial_grad_weight_ptr, + grad_weight_ptr, + num_partials, + num_cols: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +) -> None: + col_block_idx = tl.program_id(0) + partial_offsets = tl.arange(0, BLOCK_M) + col_offsets = col_block_idx * BLOCK_N + tl.arange(0, BLOCK_N) + offsets = partial_offsets[:, None] * num_cols + col_offsets[None, :] + mask = (partial_offsets[:, None] < num_partials) & (col_offsets[None, :] < num_cols) + partial = tl.load( + partial_grad_weight_ptr + offsets, + mask=mask, + other=0.0, + ) + grad_weight = tl.sum(partial, axis=0) + tl.store( + grad_weight_ptr + col_offsets, + grad_weight, + mask=col_offsets < num_cols, + ) + + +@torch.library.triton_op("torchtitan::triton_offset_rms_norm", mutates_args={}) +def _triton_offset_rms_norm_op( + input: torch.Tensor, weight: torch.Tensor, eps: float +) -> tuple[torch.Tensor, torch.Tensor]: + input = input.contiguous() + weight = weight.contiguous() + num_cols = input.shape[-1] + block_size = triton.next_power_of_2(num_cols) + if block_size > _MAX_BLOCK_SIZE: + raise ValueError( + f"Triton OffsetRMSNorm supports at most {_MAX_BLOCK_SIZE} columns, " + f"got {num_cols}" + ) + num_rows = input.numel() // num_cols + output = torch.empty_like(input) + inverse_rms = torch.empty(num_rows, dtype=torch.float32, device=input.device) + torch.library.wrap_triton(_offset_rms_norm_forward_kernel)[(num_rows,)]( + input, + weight, + output, + inverse_rms, + num_cols=num_cols, + eps=eps, + BLOCK_SIZE=block_size, + num_warps=_num_warps(block_size), + ) + return output, inverse_rms + + +@torch.library.triton_op("torchtitan::triton_offset_rms_norm_backward", mutates_args={}) +def _triton_offset_rms_norm_backward_op( + grad_output: torch.Tensor, + input: torch.Tensor, + weight: torch.Tensor, + inverse_rms: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + grad_output = grad_output.contiguous() + num_cols = input.shape[-1] + num_rows = input.numel() // num_cols + block_size = triton.next_power_of_2(num_cols) + + grad_input = torch.empty_like(input) + torch.library.wrap_triton(_offset_rms_norm_input_grad_kernel)[(num_rows,)]( + grad_output, + input, + weight, + inverse_rms, + grad_input, + num_cols=num_cols, + BLOCK_SIZE=block_size, + num_warps=_num_warps(block_size), + ) + + num_partials = triton.cdiv(num_rows, _DW_BLOCK_M) + partial_grad_weight = torch.empty( + num_partials, + num_cols, + dtype=torch.float32, + device=input.device, + ) + num_col_blocks = triton.cdiv(num_cols, _DW_BLOCK_N) + torch.library.wrap_triton(_offset_rms_norm_weight_grad_partial_kernel)[ + (num_col_blocks, num_partials) + ]( + grad_output, + input, + inverse_rms, + partial_grad_weight, + num_rows, + num_cols=num_cols, + BLOCK_M=_DW_BLOCK_M, + BLOCK_N=_DW_BLOCK_N, + num_warps=4, + ) + + grad_weight = torch.empty_like(weight) + reduce_block_m = triton.next_power_of_2(num_partials) + torch.library.wrap_triton(_offset_rms_norm_weight_grad_reduce_kernel)[ + (num_col_blocks,) + ]( + partial_grad_weight, + grad_weight, + num_partials, + num_cols=num_cols, + BLOCK_M=reduce_block_m, + BLOCK_N=_DW_BLOCK_N, + num_warps=_num_warps(reduce_block_m), + ) + return grad_input, grad_weight + + +def _triton_offset_rms_norm_setup_context(ctx, inputs, output) -> None: + input, weight, _eps = inputs + _output, inverse_rms = output + ctx.save_for_backward(input, weight, inverse_rms) + + +def _triton_offset_rms_norm_autograd_backward( + ctx, + grad_output: torch.Tensor, + _grad_inverse_rms: torch.Tensor | None, +) -> tuple[torch.Tensor, torch.Tensor, None]: + input, weight, inverse_rms = ctx.saved_tensors + grad_input, grad_weight = _triton_offset_rms_norm_backward_op( + grad_output.contiguous(), + input, + weight, + inverse_rms, + ) + return grad_input, grad_weight, None + + +_triton_offset_rms_norm_op.register_autograd( + _triton_offset_rms_norm_autograd_backward, + setup_context=_triton_offset_rms_norm_setup_context, +) + + +def triton_offset_rms_norm( + input: torch.Tensor, + weight: torch.Tensor, + eps: float, +) -> torch.Tensor: + """Compute ``(1 + weight) * rmsnorm(input)`` with Triton.""" + output, _inverse_rms = _triton_offset_rms_norm_op( + input.contiguous(), + weight.contiguous(), + eps, + ) + return output + + +class TritonOffsetRMSNorm(OffsetRMSNorm): + """Qwen3.5 OffsetRMSNorm implemented by a fused Triton kernel.""" + + @dataclass(kw_only=True, slots=True) + class Config(OffsetRMSNorm.Config): + weight_grad_sharding: SpmdType | None = None + + def __init__(self, config: Config): + super().__init__(config) + self.weight_grad_sharding = config.weight_grad_sharding + + def forward( # pyrefly: ignore[bad-param-name-override] + self, input: torch.Tensor + ) -> torch.Tensor: + if not input.is_cuda or input.dtype not in _SUPPORTED_DTYPES: + return super().forward(input) + + weight = self.weight + if isinstance(weight, DTensor): + if self.weight_grad_sharding is None: + raise AssertionError( + "DTensor weight requires a configured gradient sharding" + ) + weight = weight.to_local( + grad_placements=resolve_placements( + self.weight_grad_sharding, + weight.device_mesh, + ) + ) + return triton_offset_rms_norm(input, weight, self.eps) + + +@override( + target=OffsetRMSNorm.Config, + exact=True, + description="Use a fused Triton forward/backward kernel for Qwen3.5 OffsetRMSNorm.", +) +def triton_offset_rmsnorm( + cfg: OffsetRMSNorm.Config, +) -> TritonOffsetRMSNorm.Config: + sharding_config = cfg.sharding_config + weight_grad_sharding = None + if sharding_config is not None: + input_shardings = ( + sharding_config.in_dst_shardings or sharding_config.in_src_shardings or {} + ) + input_sharding = input_shardings.get("input") + output_sharding = ( + sharding_config.out_src_shardings or sharding_config.out_dst_shardings + ) + weight_sharding = sharding_config.state_shardings.get("weight") + if input_sharding is None or output_sharding is None: + raise ValueError( + "Triton OffsetRMSNorm requires input and output sharding " + "contracts when a sharding config is present" + ) + if weight_sharding is None: + raise ValueError("Triton OffsetRMSNorm requires a weight sharding contract") + weight_grad_sharding = SpmdType( + { + axis: axis_type.backward_type() + for axis, axis_type in weight_sharding.local_type.items() + }, + partition_spec=weight_sharding.partition_spec, + ) + sharding_config = replace( + sharding_config, + local_map=LocalMapConfig(in_grad_placements=(input_sharding,)), + ) + return derive( + cfg, + TritonOffsetRMSNorm.Config, + sharding_config=sharding_config, + weight_grad_sharding=weight_grad_sharding, + ) diff --git a/torchtitan/tools/profiler.py b/torchtitan/tools/profiler.py index 5a6bf2d1ee..bc9fcb91d6 100644 --- a/torchtitan/tools/profiler.py +++ b/torchtitan/tools/profiler.py @@ -6,6 +6,7 @@ """Kineto profiler + memory-snapshot lifecycle.""" +import inspect import os import pickle import time @@ -13,11 +14,19 @@ import torch from torchtitan.config import Configurable -from torchtitan.distributed.cudagraph import cudagraph_annotate_trace_post_processor +from torchtitan.distributed.cudagraph import get_cudagraph_annotations from torchtitan.observability import structured_logger as sl from torchtitan.tools.logging import logger from torchtitan.tools.utils import device_module +# torch's export_chrome_trace gained cuda_graph_annotations when the offline joiner +# (torch.cuda._annotate_cuda_graph_trace) was removed. Older versions still export, just +# without the CUDA graph annotations baked in. +_EXPORT_SUPPORTS_ANNOTATIONS = ( + "cuda_graph_annotations" + in inspect.signature(torch.profiler.profile.export_chrome_trace).parameters +) + # Paths expects by meta internal tooling PROFILE_DIR = "profiling/traces" # Profiler.Config.save_traces_folder default PROFILE_ITER_DIR = "iteration_{step}" # PROFILE_DIR/{PROFILE_ITER_DIR} @@ -295,8 +304,19 @@ def trace_handler(prof): begin = time.monotonic() output_file = os.path.join(curr_trace_dir, PROFILE_FILE.format(rank=rank)) - prof.export_chrome_trace(output_file) - cudagraph_annotate_trace_post_processor(output_file) + # CUDA graph annotations are baked in during the export rather than + # joined onto the written file afterwards: re-reading and rewriting a + # gzipped trace paid the compression cost twice. + annotations = get_cudagraph_annotations() + if annotations and not _EXPORT_SUPPORTS_ANNOTATIONS: + logger.warning( + "This torch does not support cuda_graph_annotations on " + "export_chrome_trace; the trace will have no CUDA graph kernel " + "annotations." + ) + annotations = None + extra = {"cuda_graph_annotations": annotations} if annotations else {} + prof.export_chrome_trace(output_file, **extra) logger.info( f"Finished dumping profiler traces in {time.monotonic() - begin:.2f} seconds" diff --git a/torchtitan/tools/utils.py b/torchtitan/tools/utils.py index 5fffd8ae33..614cbf0426 100644 --- a/torchtitan/tools/utils.py +++ b/torchtitan/tools/utils.py @@ -35,7 +35,10 @@ def has_cuda_capability(major: int, minor: int) -> bool: def get_cuda_flash_attention_impl() -> str | None: """Return the FlashAttention implementation for the current CUDA architecture.""" - # Blackwell (SM 10.0) and newer use FA4; Hopper (SM 9.0) uses FA3. + + # FA4 advertises Hopper support, but as of writing it hangs under + # torch.compile there, so Hopper (sm90) stays on FA3. + # https://github.com/pytorch/torchtitan/pull/4413 if has_cuda_capability(10, 0): return "FA4" if has_cuda_capability(9, 0): diff --git a/torchtitan_recipes/tests/b200.py b/torchtitan_recipes/tests/b200.py index 38591612c6..e445e1cfd1 100644 --- a/torchtitan_recipes/tests/b200.py +++ b/torchtitan_recipes/tests/b200.py @@ -15,3 +15,11 @@ def kimi_k3_debugmodel_mm_fsdp2() -> Trainer.Config: config = kimi_k3_debugmodel() config.parallelism.data_parallel_shard_degree = 2 return config + + +def llama3_debugmodel_mxfp8_fsdp2() -> Trainer.Config: + from torchtitan.models.llama3.config_registry import llama3_debugmodel_mxfp8 + + config = llama3_debugmodel_mxfp8() + config.parallelism.data_parallel_shard_degree = 2 + return config