[https://nvbugs/5463720][fix] tp-split the inferred mlp_hidden_size for nemotron-nas (NVIDIA#7231)

Signed-off-by: Venky Ganesh <[email protected]>
Signed-off-by: Wangshanshan <[email protected]>

Author: venkywonka

tensorrt_llm/_torch/model_config.py

@@ -494,7 +494,8 @@ def get_bindings_model_config(self,
                 architectures = self.pretrained_config.architectures
                 if len(architectures
                        ) == 1 and architectures[0] == "DeciLMForCausalLM":
-                    mlp_hidden_size = self._infer_nemotron_ffn_mult()
+                    mlp_hidden_size = self._infer_nemotron_ffn_mult(
+                    ) // self.mapping.tp_size
                 else:
                     raise ValueError(
                         f"Inferring mlp hidden size for model architecture: {architectures} isn't supported yet"