From 53ffe460eb82dba9a893be03557e4b639456ea5f Mon Sep 17 00:00:00 2001 From: Chelsea Lowman Date: Tue, 26 May 2026 16:11:48 +0000 Subject: [PATCH] Updating gpt-oss-20b to use mxfp8 --- torchtitan/models/gpt_oss/config_registry.py | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/torchtitan/models/gpt_oss/config_registry.py b/torchtitan/models/gpt_oss/config_registry.py index b3fef6355e..abde4e6eb2 100644 --- a/torchtitan/models/gpt_oss/config_registry.py +++ b/torchtitan/models/gpt_oss/config_registry.py @@ -4,7 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -from torchtitan.components.quantization import Float8LinearConverter, Float8GroupedExpertsConverter +from torchtitan.components.quantization import Float8LinearConverter, Float8GroupedExpertsConverter, MXFP8LinearConverter, MXFP8GroupedExpertsConverter from torchtitan.components.checkpoint import CheckpointManager from torchtitan.components.loss import ChunkedCELoss, CrossEntropyLoss from torchtitan.components.lr_scheduler import LRSchedulersContainer @@ -25,7 +25,7 @@ def gpt_oss_debugmodel() -> Trainer.Config: return Trainer.Config( - loss=ChunkedCELoss.Config(), + loss=CrossEntropyLoss.Config(), hf_assets_path="./tests/assets/tokenizer", metrics=MetricsProcessor.Config(log_freq=1), model_spec=model_registry("debugmodel"), @@ -72,13 +72,13 @@ def gpt_oss_20b() -> Trainer.Config: loss=CrossEntropyLoss.Config(), hf_assets_path="./assets/hf/gpt-oss-20b", model_spec=model_registry("20b", converters=[ - Float8LinearConverter.Config( + MXFP8LinearConverter.Config( model_compile_enabled=True, - filter_fqns=["output", "router.gate"], + fqns=["qkv_linear", "wo"], + ), + MXFP8GroupedExpertsConverter.Config( + model_compile_enabled=True ), - Float8GroupedExpertsConverter.Config( - model_compile_enabled=True, - ) ]), dataloader=HuggingFaceTextDataLoader.Config(dataset="c4", pin_memory=True, num_workers=1, prefetch_factor=2), optimizer=OptimizersContainer.Config(lr=8e-4), @@ -89,15 +89,16 @@ def gpt_oss_20b() -> Trainer.Config: min_lr_factor=0.1, ), training=TrainingConfig( - local_batch_size=2, + local_batch_size=8, seq_len=8192, steps=10000, ), parallelism=ParallelismConfig( expert_parallel_degree=1, + fsdp_reshard_after_forward="never", ), checkpoint=CheckpointManager.Config(interval=500), - activation_checkpoint=ActivationCheckpointConfig(mode="full"), + activation_checkpoint=ActivationCheckpointConfig(mode="selective"), compile=CompileConfig(enable=True), )