From cb6347c956d0f60e371c4ed292b9051ec6033af8 Mon Sep 17 00:00:00 2001 From: Stanley Sun Date: Tue, 27 Oct 2026 23:38:12 -0700 Subject: [PATCH] [None][test] Add gap-fixer coverage for job 107 Signed-off-by: Stanley Sun --- .../_torch/modeling/test_modeling_bart.py | 88 ++++++++++++++++++- 1 file changed, 86 insertions(+), 2 deletions(-) diff --git a/tests/unittest/_torch/modeling/test_modeling_bart.py b/tests/unittest/_torch/modeling/test_modeling_bart.py index 0ec2d9d778e4..aa389672d861 100644 --- a/tests/unittest/_torch/modeling/test_modeling_bart.py +++ b/tests/unittest/_torch/modeling/test_modeling_bart.py @@ -17,13 +17,17 @@ import pytest import torch -from transformers import BartConfig +from transformers import BartConfig, MBartConfig from transformers import BartForConditionalGeneration as HFBartForConditionalGeneration +from transformers import MBartForConditionalGeneration as HFMBartForConditionalGeneration from tensorrt_llm._torch.attention.backends.vanilla import VanillaAttentionMetadata from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm._torch.models.modeling_bart import BartForConditionalGeneration +from tensorrt_llm._torch.models.modeling_bart import ( + BartForConditionalGeneration, + MBartForConditionalGeneration, +) from tensorrt_llm._torch.models.modeling_utils import get_registered_model_class @@ -111,3 +115,83 @@ def test_bart_for_conditional_generation_construction_and_forward() -> None: # Both implementations use float32 and identical small weights. Allow # accumulation-order differences between SDPA and eager HF attention. torch.testing.assert_close(logits, expected, rtol=2e-4, atol=2e-5) + + +@pytest.mark.skipif( + not torch.cuda.is_available(), reason="The production mBART layers require CUDA" +) +def test_mbart_for_conditional_generation_construction_and_forward() -> None: + torch.manual_seed(0) + config = MBartConfig( + vocab_size=64, + d_model=32, + encoder_layers=1, + decoder_layers=1, + encoder_attention_heads=4, + decoder_attention_heads=4, + encoder_ffn_dim=64, + decoder_ffn_dim=64, + max_position_embeddings=32, + activation_function="relu", + dropout=0.0, + attention_dropout=0.0, + activation_dropout=0.0, + scale_embedding=True, + tie_word_embeddings=True, + pad_token_id=1, + bos_token_id=0, + eos_token_id=2, + decoder_start_token_id=2, + forced_eos_token_id=2, + torch_dtype=torch.float32, + ) + config._attn_implementation = "eager" + reference = HFMBartForConditionalGeneration(config).cuda().float().eval() + model_cls = get_registered_model_class("MBartForConditionalGeneration") + assert model_cls is MBartForConditionalGeneration + model = model_cls(ModelConfig(pretrained_config=config, attn_backend="VANILLA")).cuda().eval() + assert isinstance(model, MBartForConditionalGeneration) + assert model.model.encoder.layer_norm is not None + assert model.model.decoder.layer_norm is not None + assert model.model.encoder.layers[0].normalize_before + assert model.model.decoder.layers[0].normalize_before + model.load_weights( + {name: value.detach().clone() for name, value in reference.state_dict().items()} + ) + + encoder_input_ids = torch.tensor([3, 7, 11, 5, 9], dtype=torch.int32, device="cuda") + decoder_input_ids = torch.tensor([2, 4, 6], dtype=torch.int32, device="cuda") + encoder_metadata = _context_metadata(encoder_input_ids.numel()) + decoder_metadata = _context_metadata(decoder_input_ids.numel()) + cross_metadata = decoder_metadata.create_cross_metadata([encoder_input_ids.numel()]) + cross_metadata.prepare() + assert cross_metadata.is_cross + + # Learned position tables reserve their first two entries. + encoder_positions = ( + torch.arange(encoder_input_ids.numel(), dtype=torch.int32, device="cuda") + 2 + ) + decoder_positions = ( + torch.arange(decoder_input_ids.numel(), dtype=torch.int32, device="cuda") + 2 + ) + with torch.inference_mode(): + logits = model.forward( + attn_metadata=decoder_metadata, + input_ids=decoder_input_ids, + position_ids=decoder_positions, + encoder_input_ids=encoder_input_ids, + encoder_position_ids=encoder_positions, + encoder_attn_metadata=encoder_metadata, + cross_attn_metadata=cross_metadata, + return_context_logits=True, + ) + expected = reference( + input_ids=encoder_input_ids.long().unsqueeze(0), + decoder_input_ids=decoder_input_ids.long().unsqueeze(0), + use_cache=False, + ).logits.squeeze(0) + + assert logits.shape == (decoder_input_ids.numel(), config.vocab_size) + assert logits.dtype == torch.float32 + assert torch.isfinite(logits).all() + torch.testing.assert_close(logits, expected, rtol=2e-4, atol=2e-5)