fix: Support TLLM_OVERRIDE_LAYER_NUM for llama4. (#3679)

Signed-off-by: Yuxian Qiu <[email protected]>

Author: yuxianq

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -771,6 +771,10 @@ def _load_model(self, checkpoint_dir: str, load_format: LoadFormat,
         num_layers = int(os.environ.get("TLLM_OVERRIDE_LAYER_NUM", "0"))
         if num_layers > 0:
             config.pretrained_config.num_hidden_layers = num_layers
+            for sub_config in ["text_config", "vision_config"]:
+                if hasattr(config.pretrained_config, sub_config):
+                    getattr(config.pretrained_config,
+                            sub_config).num_hidden_layers = num_layers
 
         with timing("Model init total"):
             try: