fix gptoss rope config (#1445)

Author: hhaAndroid

tests/model/test_gpt_oss_moe.py

@@ -13,7 +13,6 @@
 from xtuner.v1.model.moe.moe import SequenceContext
 from xtuner.v1.model.moe.gpt_oss import GptOss21BA3P6Config
 from xtuner.v1.config import FSDPConfig
-from xtuner.v1.utils.compile import maybe_compile
 from xtuner.v1.loss.ce_loss import CELossConfig, CELossContextInputItem
 
 GPT_OSS_MINI_PATH = os.environ["GPT_OSS_MINI_PATH"]
@@ -45,7 +44,6 @@ def test_gpt_oss_run(self, device, dispatcher, ep_size, compile, tol, loss_class
         self.create_pg(device)
 
         hf_config = AutoConfig.from_pretrained(GPT_OSS_MINI_PATH)
-        hf_config.rope_scaling = None
 
         hf_model = AutoModelForCausalLM.from_pretrained(
             GPT_OSS_MINI_PATH,
@@ -108,7 +106,6 @@ def test_fsdp_accuracy(self, device, dispatcher, ep_size):
         self.create_pg(device)
 
         hf_config = AutoConfig.from_pretrained(GPT_OSS_MINI_PATH)
-        hf_config.rope_scaling = None
         hf_model = AutoModelForCausalLM.from_pretrained(
             GPT_OSS_MINI_PATH,
             dtype=torch.bfloat16,

xtuner/v1/model/base.py

@@ -124,6 +124,13 @@ class TransformerConfig(XTunerBaseModelConfig):
         "default"
     )
 
+    @computed_field  # type: ignore[misc]
+    @property
+    def rope_scaling(self) -> dict | None:
+        if self.rope_scaling_cfg is not None:
+            return self.rope_scaling_cfg.model_dump()
+        return None
+
     @computed_field
     def num_attention_heads(self) -> int:
         return self.attention.num_attention_heads

xtuner/v1/model/moe/gpt_oss.py

@@ -10,6 +10,7 @@
 from xtuner.v1.model.moe.moe import BalancingLossConfig, MoEConfig
 from xtuner.v1.module.attention import MHAConfig
 from xtuner.v1.module.decoder_layer.moe_decoder_layer import MoEActFnConfig
+from xtuner.v1.module.rope import RopeScalingConfig
 from xtuner.v1.module.router.greedy import GreedyRouterConfig
 
 from .moe import MoE
@@ -123,6 +124,9 @@ class GptOssConfig(MoEConfig):
     tie_word_embeddings: bool = False
     n_shared_experts: int = 0
     moe_act_fn_cfg: MoEActFnConfig = MoEActFnConfig(act_type="clipped_swiglu", clip_alpha=1.702, clip_limit=7)
+    rope_scaling_cfg: RopeScalingConfig = RopeScalingConfig(
+        type="yarn", beta_fast=32.0, beta_slow=1.0, factor=32.0, original_max_position_embeddings=4096, truncate=False
+    )
 
     @computed_field
     def layers_type(self) -> list[Literal["full_attention", "sliding_attention"]]:
@@ -138,7 +142,6 @@ def from_hf(cls, hf_path: str | Path) -> Self:
         assert isinstance(cfg, HFGptOssConfig)
 
         config = cls(
-            hf_config=cfg,
             vocab_size=cfg.vocab_size,
             max_position_embeddings=cfg.max_position_embeddings,
             pad_token_id=cfg.pad_token_id,
@@ -168,8 +171,17 @@ def from_hf(cls, hf_path: str | Path) -> Self:
                 norm_topk_prob=True,
                 router_scaling_factor=1.0,
             ),
+            rope_scaling_cfg=RopeScalingConfig(
+                type=cfg.rope_scaling.get("rope_type", "yarn"),
+                beta_fast=cfg.rope_scaling.get("beta_fast", 32.0),
+                beta_slow=cfg.rope_scaling.get("beta_slow", 1.0),
+                factor=cfg.rope_scaling.get("factor", 32.0),
+                original_max_position_embeddings=cfg.rope_scaling.get("original_max_position_embeddings", 4096),
+                truncate=cfg.rope_scaling.get("truncate", False),
+            )
+            if cfg.rope_scaling is not None
+            else None,
         )
-
         return config
 
     @property
@@ -201,6 +213,16 @@ def hf_config(self) -> HFGptOssConfig:
             o_bias=True,
             dtype=torch.bfloat16,
             swiglu_limit=self.moe_act_fn_cfg.clip_limit,
+            rope_scaling={
+                "rope_type": self.rope_scaling_cfg.type,
+                "beta_fast": self.rope_scaling_cfg.beta_fast,
+                "beta_slow": self.rope_scaling_cfg.beta_slow,
+                "factor": self.rope_scaling_cfg.factor,
+                "original_max_position_embeddings": self.rope_scaling_cfg.original_max_position_embeddings,
+                "truncate": self.rope_scaling_cfg.truncate,
+            }
+            if self.rope_scaling_cfg is not None
+            else None,
         )
 
 

xtuner/v1/module/rope/rope.py

@@ -20,13 +20,12 @@ class RopeScalingConfig(BaseModel):
     model_config = ConfigDict(extra="forbid")
     type: Literal["default", "linear", "dynamic", "yarn", "longrope", "llama3", "qwen3_vl"] = "default"
 
-    max_position_embeddings: int | None = None  # TODO: 无用参数考虑删除
-    original_max_position_embeddings: int | None = None  # TODO: 无用参数考虑删除
+    max_position_embeddings: int | None = None
+    original_max_position_embeddings: int | None = None
 
     # For Qwen3VL
     mrope_section: list[int] | None = None  # e.g. [24, 20, 20]
 
-    # For inference
     factor: float | None = None
     beta_fast: float | None = None
     beta_slow: float | None = None
@@ -36,6 +35,7 @@ class RopeScalingConfig(BaseModel):
     high_freq_factor: float | None = None
     mscale: float | None = None
     mscale_all_dim: float | None = None
+    truncate: bool = False
 
     # For FoPE
     fope_init_factor: float | None = None
@@ -73,6 +73,14 @@ def __init__(self, config, device=None):
         self.original_max_seq_len = config.max_position_embeddings
         self.rope_type = "default"
         self.config = config
+
+        rope_scaling_cfg = config.rope_scaling_cfg
+        if rope_scaling_cfg is not None:
+            self.rope_type = rope_scaling_cfg.type
+        assert self.rope_type in ["default", "linear", "yarn", "llama3"], (
+            f"Unsupported rope_type: {self.rope_type}. Supported types are: 'default', 'linear', 'yarn', 'llama3'."
+        )
+
         self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
 
         inv_freq: torch.Tensor