[https://nvbugs/5552132][fix] Enable LoRa for GPT OSS Torch

tensorrt_llm/_torch/models/modeling_gpt_oss.py

@@ -230,21 +230,27 @@ def _create_ideal_expert_load_balanced_logits(
             device=device,
             dtype=pretrained_config.torch_dtype)
 
-    def compute_gate_output(self, x: torch.Tensor) -> torch.Tensor:
-        if get_sm_version() in [
-                90, 100, 103
-        ] and x.shape[0] <= MIN_LATENCY_TINYGEMM_NUM_TOKENS:
+    def compute_gate_output(self,
+                            x: torch.Tensor,
+                            lora_params: Optional[dict] = None) -> torch.Tensor:
+        # Skip tinygemm2 optimization when LoRA is active (tinygemm2 doesn't support LoRA)
+        use_tinygemm = (get_sm_version() in [90, 100, 103]
+                        and x.shape[0] <= MIN_LATENCY_TINYGEMM_NUM_TOKENS
+                        and (lora_params is None or not bool(lora_params)))
+
+        if use_tinygemm:
             weight = self.gate.weight
             bias = self.gate.bias
             g = torch.ops.trtllm.tinygemm2(x, weight, bias)
         else:
-            g = self.gate(x)
+            g = self.gate(x, lora_params=lora_params, layer_idx=self.layer_idx)
         return g
 
     def forward_normal(
         self,
         x: torch.Tensor,
-        residual: Optional[torch.Tensor] = None
+        residual: Optional[torch.Tensor] = None,
+        lora_params: Optional[dict] = None
     ) -> tuple[torch.Tensor, torch.Tensor]:
         orig_shape = x.shape
         hidden_dim = orig_shape[-1]
@@ -253,7 +259,7 @@ def forward_normal(
         # t = self.norm(x) was done in the parent block
         t = x
 
-        g = self.compute_gate_output(t)
+        g = self.compute_gate_output(t, lora_params=lora_params)
         # Use ideal load balanced logits if enabled, otherwise use gate output
         if os.environ.get('ENABLE_PERFECT_ROUTER', '0') == '1':
             # WARNING: This discards the learned gate output and uses ideal logits for perfect load balancing
@@ -272,7 +278,8 @@ def forward_attn_dp(
         self,
         x: torch.Tensor,
         attn_metadata: AttentionMetadata,
-        residual: Optional[torch.Tensor] = None
+        residual: Optional[torch.Tensor] = None,
+        lora_params: Optional[dict] = None
     ) -> tuple[torch.Tensor, torch.Tensor]:
         orig_shape = x.shape
         hidden_dim = orig_shape[-1]
@@ -288,7 +295,7 @@ def forward_attn_dp(
             if (isinstance(self.experts, (TritonFusedMoE))):
                 t = allgather(t, self.mapping, dim=0, sizes=all_rank_num_tokens)
 
-        g = self.compute_gate_output(t)
+        g = self.compute_gate_output(t, lora_params=lora_params)
         # Use ideal load balanced logits if enabled, otherwise use gate output
         if os.environ.get('ENABLE_PERFECT_ROUTER', '0') == '1':
             # WARNING: This discards the learned gate output and uses ideal logits for perfect load balancing
@@ -312,12 +319,13 @@ def forward(
         self,
         x: torch.Tensor,
         attn_metadata: AttentionMetadata,
-        residual: Optional[torch.Tensor] = None
+        residual: Optional[torch.Tensor] = None,
+        lora_params: Optional[dict] = None
     ) -> tuple[torch.Tensor, torch.Tensor]:
         if self.enable_attention_dp:
-            return self.forward_attn_dp(x, attn_metadata, residual)
+            return self.forward_attn_dp(x, attn_metadata, residual, lora_params)
         else:
-            return self.forward_normal(x, residual)
+            return self.forward_normal(x, residual, lora_params)
 
 
 class TransformerBlock(DecoderLayer):
@@ -378,6 +386,7 @@ def forward_normal(
         attn_metadata: AttentionMetadata,
         residual: Optional[torch.Tensor] = ...,
         spec_metadata: Optional[SpecMetadata] = None,
+        lora_params: Optional[dict] = None,
         **kwargs,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         if residual is None:
@@ -388,10 +397,14 @@ def forward_normal(
                                 hidden_states,
                                 attn_metadata,
                                 residual=residual,
+                                lora_params=lora_params,
                                 **kwargs)
         x, residual = self.post_attention_layernorm(x, residual)
 
-        x, residual = self.mlp(x, attn_metadata, residual)
+        x, residual = self.mlp(x,
+                               attn_metadata,
+                               residual,
+                               lora_params=lora_params)
 
         if spec_metadata is not None:
             spec_metadata.maybe_capture_hidden_states(self.layer_idx, x,
@@ -407,6 +420,7 @@ def forward_tp(
         attn_metadata: AttentionMetadata,
         residual: Optional[torch.Tensor] = ...,
         spec_metadata: Optional[SpecMetadata] = None,
+        lora_params: Optional[dict] = None,
         **kwargs,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         if residual is None:
@@ -419,6 +433,7 @@ def forward_tp(
             attn_metadata,
             residual=residual,
             all_reduce_params=AllReduceParams(enable_allreduce=False),
+            lora_params=lora_params,
             **kwargs)
 
         x, residual = self.allreduce(
@@ -431,7 +446,10 @@ def forward_tp(
                 trigger_completion_at_end=False,
             ))
 
-        x, residual = self.mlp(x, attn_metadata, residual)
+        x, residual = self.mlp(x,
+                               attn_metadata,
+                               residual,
+                               lora_params=lora_params)
 
         if spec_metadata is not None and spec_metadata.is_layer_capture(
                 self.layer_idx):
@@ -467,6 +485,7 @@ def forward(
         attn_metadata: AttentionMetadata,
         residual: Optional[torch.Tensor] = ...,
         spec_metadata: Optional[SpecMetadata] = None,
+        lora_params: Optional[dict] = None,
         **kwargs,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         if self.is_tp:
@@ -478,6 +497,7 @@ def forward(
                         attn_metadata,
                         residual,
                         spec_metadata=spec_metadata,
+                        lora_params=lora_params,
                         **kwargs)
 
 
@@ -540,6 +560,7 @@ def forward(
         position_ids: Optional[torch.LongTensor] = None,
         inputs_embeds: Optional[torch.FloatTensor] = None,
         spec_metadata: Optional[SpecMetadata] = None,
+        lora_params: Optional[dict] = None,
         **kwargs,
     ) -> torch.Tensor:
         if (input_ids is None) ^ (inputs_embeds is not None):
@@ -557,6 +578,7 @@ def forward(
                 attn_metadata=attn_metadata,
                 residual=residual,
                 spec_metadata=spec_metadata,
+                lora_params=lora_params,
             )
 
         return hidden_states

tests/integration/defs/conftest.py

@@ -631,6 +631,15 @@ def deepseek_v3_model_root(request):
     return deepseek_v3_model_root
 
 
+@pytest.fixture(scope="function")
+def gpt_oss_model_root(request):
+    models_root = os.path.join(llm_models_root(), "gpt_oss")
+    if request.param == "gpt-oss-20b":
+        gpt_oss_model_root = os.path.join(models_root, "gpt-oss-20b")
+    assert exists(gpt_oss_model_root), f"{gpt_oss_model_root} does not exist!"
+    return gpt_oss_model_root
+
+
 @pytest.fixture(scope="session")
 def trt_performance_cache_name():
     return "performance.cache"
@@ -1345,6 +1354,10 @@ def llm_lora_model_root(request):
                 os.path.join(
                     models_root, "nemotron-nas",
                     "Llama-3_3-Nemotron-Super-49B-v1-lora-adapter_NIM_r32"))
+        elif item == "gpt-oss-20b-lora-adapter_NIM_r8":
+            model_root_list.append(
+                os.path.join(models_root, "gpt_oss",
+                             "gpt-oss-20b-lora-adapter_NIM_r8"))
 
     return ",".join(model_root_list)
 

tests/integration/defs/examples/test_gpt.py

@@ -18,6 +18,7 @@
 import re
 from pathlib import Path
 
+import defs.ci_profiler
 import pytest
 from defs.common import (convert_weights, generate_summary_cmd, parse_mpi_cmd,
                          parse_output, quantize_data, run_and_check, similar,
@@ -28,6 +29,11 @@
                            skip_post_blackwell, skip_pre_ada)
 from defs.trt_test_alternative import check_call
 
+from tensorrt_llm import LLM
+from tensorrt_llm.executor.request import LoRARequest
+from tensorrt_llm.lora_manager import LoraConfig
+from tensorrt_llm.sampling_params import SamplingParams
+
 # skip trt flow cases on post-Blackwell-Ultra
 if get_sm_version() >= 103:
     pytest.skip(
@@ -1928,3 +1934,62 @@ def test_llm_minitron_fp8_with_pseudo_loras(gpt_example_root,
         target_trtllm_modules=["attn_q", "attn_k", "attn_v"],
         zero_lora_weights=True,
     )
+
+
+@pytest.mark.skip_less_device_memory(
+    20000)  # Conservative 20GB requirement for GPT-OSS-20B
+@pytest.mark.parametrize("gpt_oss_model_root", [
+    "gpt-oss-20b",
+], indirect=True)
+@pytest.mark.parametrize("llm_lora_model_root",
+                         ['gpt-oss-20b-lora-adapter_NIM_r8'],
+                         indirect=True)
+def test_gpt_oss_20b_lora_torch(gpt_example_root, llm_venv, gpt_oss_model_root,
+                                llm_datasets_root, llm_rouge_root, engine_dir,
+                                cmodel_dir, llm_lora_model_root):
+    """Run GPT-OSS-20B with LoRA adapter using Torch backend."""
+
+    print(f"Using LoRA from: {llm_lora_model_root}")
+
+    defs.ci_profiler.start("test_gpt_oss_20b_lora_torch")
+
+    lora_config = LoraConfig(
+        lora_dir=[llm_lora_model_root],
+        max_lora_rank=8,  # Match adapter_config.json "r": 8
+        max_loras=1,
+        max_cpu_loras=1,
+    )
+
+    with LLM(model=gpt_oss_model_root, lora_config=lora_config) as llm:
+
+        prompts = [
+            "User: Message Mason saying that we should compete in next week's football tournament, and tell him that the winner will get $100.\n\nAssistant: "
+        ]
+
+        sampling_params = SamplingParams(max_tokens=50)
+
+        lora_request = [LoRARequest("gpt-oss-lora", 0, llm_lora_model_root)]
+
+        print("Running inference with real LoRA adapter...")
+        outputs = llm.generate(prompts,
+                               sampling_params,
+                               lora_request=lora_request)
+
+        expected_output = " Hey Mason! I hope you're doing well. I was thinking about the next week's football tournament and I wanted to give you a hint that we should compete in it. The winner will be a great opportunity for us to win $100.\n\nUser:"
+
+        for i, output in enumerate(outputs):
+            print(f"Prompt {i+1}: {prompts[i]}")
+            print(f"Response {i+1}: {output.outputs[0].text}")
+            print("-" * 50)
+
+        assert len(outputs) == 1
+        assert len(outputs[0].outputs) > 0
+        generated_text = outputs[0].outputs[0].text
+        similarity = similarity_score(generated_text, expected_output)
+        assert similar(generated_text, expected_output, threshold=0.8), \
+            f"Output similarity too low (similarity={similarity:.2%})!\nExpected: {repr(expected_output)}\nGot: {repr(generated_text)}"
+
+    defs.ci_profiler.stop("test_gpt_oss_20b_lora_torch")
+    print(
+        f"test_gpt_oss_20b_lora_torch: {defs.ci_profiler.elapsed_time_in_sec('test_gpt_oss_20b_lora_torch')} sec"
+    )

tests/integration/test_lists/test-db/l0_h100.yml

@@ -395,6 +395,7 @@ l0_h100:
   - examples/test_phi.py::test_llm_phi_quantization_1gpu[Phi-4-mini-instruct-fp8-bfloat16]
   - unittest/trt/model/test_gpt_e2e.py # 3 mins / 6 mins on H100
   - unittest/trt/attention/test_gpt_attention_no_cache.py
+  - examples/test_gpt.py::test_gpt_oss_20b_lora_torch[gpt-oss-20b-lora-adapter_NIM_r8-gpt-oss-20b]
 - condition:
     ranges:
       system_gpu_count: