tensorrt_llm/_torch/model_config.py (3)

244-247: Enable rope_scaling validation early.

You define _rope_scaling_validation() but never call it. Enabling it in __init__ will catch invalid configs early.

Apply this diff:

         self.rope_theta = rope_theta
         self.rope_scaling = rope_scaling
-        # self._rope_scaling_validation()   # TODO: Need validation?
+        self._rope_scaling_validation()

If strictness is a concern, you can guard this call behind an env flag or parameter.

---

34-116: Docstring defaults and descriptions are out of sync with the implementation.

Examples:

• vocab_size doc says 32000, code defaults to 290943.
• rms_norm_eps doc says 1e-06, code uses 1e-5.
• moe_intermediate_size doc says defaults to 11008, code defaults to None.
• attention_bias has duplicated “defaults to False”.

Please align the docstring with the actual defaults to avoid confusion for users.

---

352-383: Be less strict about rope_scaling keys.

The current validation requires the dict to have exactly two keys (len(...) != 2 check). In practice, configs may carry auxiliary fields; this strictness can cause unnecessary breakage. Consider validating presence and type of required keys without constraining the total number of keys.

tensorrt_llm/_torch/models/modeling_hunyuan_dense.py (5)

9-12: Use project logger instead of print and add logger import.

This module emits print statements later. Prefer the project logger for consistency and log-level control.

Apply this diff to import the logger:

 from tqdm import tqdm
 
 from tensorrt_llm._torch.distributed import AllReduceParams
 from tensorrt_llm.functional import PositionEmbeddingType
+from tensorrt_llm.logger import logger

And see further comments for replacing print calls below.

---

250-255: Consider using HunYuanMLP instead of the generic GatedMLP.

You introduced HunYuanMLP to encapsulate HunYuan-specific sizing and MoE/shared-MLP choices, but the layer currently instantiates GatedMLP directly. Replace with HunYuanMLP to honor moe_intermediate_size, num_shared_expert, etc.

Apply this diff:

-        self.mlp = GatedMLP(hidden_size=config.hidden_size,
-                            intermediate_size=config.intermediate_size,
-                            bias=config.mlp_bias,
-                            dtype=config.torch_dtype,
-                            config=model_config)
+        self.mlp = HunYuanMLP(
+            model_config=model_config,
+            layer_idx=layer_idx,
+            is_shared_mlp=False,
+            is_moe=(getattr(config, "num_experts", 1) != 1),
+        )

If shared-MLP is part of your design, set is_shared_mlp=True based on config.

---

444-456: Avoid hard-coded internal paths in debug dump directory.

Defaulting TRT_MODEL_DUMP_DIR to an internal path leaks internal infra details. Use a safe local default (e.g., ./trt_model_dumps or /tmp/trt_model_dumps).

Apply this diff:

                 self.dump_modules_file = os.path.join(
                     os.environ.get(
-                        "TRT_MODEL_DUMP_DIR",
-                        "/apdcephfs_jn/share_302216743/jiarunliu/trt_model_dumps"
+                        "TRT_MODEL_DUMP_DIR",
+                        os.path.join(os.getcwd(), "trt_model_dumps"),
                     ),
                     f"trt_modules_{time.strftime('%Y%m%d_%H%M%S')}_{id(self):x}.jsonl"
                 )

---

572-573: Use logger and English messages for debug prints.

Replace print with logger.info and avoid non-English logs to keep consistent logging conventions and CI log readability.

Apply this diff:

-                    print(f"已完成第{self.decode_step - 3}次decode的数据收集，hook已移除")
+                    logger.info("Removed debug hooks after collecting decode step %d", self.decode_step - 3)
...
-                print(
-                    f"第{self.decode_step - 2}次decode的模块信息已dump到: {self.dump_modules_file}"
-                )
+                logger.info(
+                    "Dumped module info for decode step %d to %s",
+                    self.decode_step - 2,
+                    self.dump_modules_file,
+                )

Also applies to: 593-595

---

29-73: Localize or translate comments/docstrings.

Function docstrings and inline comments here are partly in Chinese. For consistency and maintainability across the project, translate them to English (Google-style docstrings preferred).

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.