From 321f72cd3cd0af0d55c9014263a800cd24afb693 Mon Sep 17 00:00:00 2001 From: Omer Ullman Argov <118735753+omera-nv@users.noreply.github.com> Date: Thu, 8 May 2025 07:03:34 +0000 Subject: [PATCH 1/9] replace sanity test for nemotron h with a correctness test Signed-off-by: Omer Ullman Argov <118735753+omera-nv@users.noreply.github.com> --- .../modeling/test_modeling_nemotron_h.py | 230 +++++++++++++++--- 1 file changed, 199 insertions(+), 31 deletions(-) diff --git a/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py b/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py index 7cec137abfd0..eff77b338a88 100644 --- a/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py +++ b/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py @@ -13,6 +13,10 @@ NemotronHForCausalLM ) # isort: on +from transformers import AutoTokenizer +from utils.llm_data import llm_models_root + +from tensorrt_llm._torch.pyexecutor.model_engine import load_weights from tensorrt_llm._torch.pyexecutor.resource_manager import \ MambaHybridCacheManager from tensorrt_llm.bindings.executor import KvCacheConfig @@ -64,29 +68,41 @@ } +def get_logprobs(token_ids: torch.Tensor, logits: torch.Tensor) -> torch.Tensor: + raw_probs = torch.softmax(logits, dim=-1) + index = token_ids.unsqueeze(1).cuda() + token_probs = torch.gather(raw_probs, dim=1, index=index).squeeze(-1) + return torch.log(token_probs) + + class TestNemotronH(unittest.TestCase): - def test_nemotron_h_sanity(self): + def test_nemotron_correctness(self): config_dict = deepcopy(NEMOTRON_H_CONFIG) nemotron_h_config = NemotronHConfig.from_dict(config_dict) + model_dir = f"{llm_models_root(check=True)}/Nemotron-H-8B-Base-8K" + + tokenizer = AutoTokenizer.from_pretrained(model_dir) + dtype = nemotron_h_config.torch_dtype device = torch.device('cuda') model_config = ModelConfig(pretrained_config=nemotron_h_config) nemotron_h = NemotronHForCausalLM(model_config).to(device) - input_ids = torch.tensor([100, 200, 300, 100, 200, 100, 400, 500], - dtype=torch.int, + weights = load_weights(model_dir) + nemotron_h.load_weights(weights) + + # These tokens are from the sentence: The future of AI is + input_ids = torch.tensor([1784, 7147, 1307, 26554, 1395], + dtype=torch.int64, device=device) - context_sequence_lengths = [3, 2, 1] - sequence_lengths = context_sequence_lengths + [1, 1] - past_seen_tokens = [0, 0, 0, 62, 75] - request_ids = list(range(len(sequence_lengths))) - token_nums = (torch.tensor(past_seen_tokens) + - torch.tensor(sequence_lengths)).tolist() - prompt_lens = token_nums[:3] + past_seen_tokens[3:] + num_cached_tokens_per_seq = [0] + request_ids = [1] + token_nums = [input_ids.size(-1)] + prompt_lens = [input_ids.size(-1)] num_blocks = 100 tokens_per_block = 128 @@ -95,7 +111,7 @@ def test_nemotron_h_sanity(self): mamba_num_layers = nemotron_h.config.hybrid_override_pattern.count("M") num_kv_heads = nemotron_h.config.num_key_value_heads max_seq_len = num_blocks * tokens_per_block - batch_size = len(context_sequence_lengths) + 2 + max_batch_size = 1 if dtype == torch.bfloat16: kv_cache_dtype = tensorrt_llm.bindings.DataType.BF16 @@ -124,7 +140,7 @@ def test_nemotron_h_sanity(self): head_dim=head_dim, tokens_per_block=tokens_per_block, max_seq_len=max_seq_len, - max_batch_size=batch_size, + max_batch_size=max_batch_size, mapping=mapping, dtype=kv_cache_dtype, num_extra_kv_tokens=0, @@ -134,44 +150,196 @@ def test_nemotron_h_sanity(self): metadata_cls = get_attention_backend(model_config.attn_backend).Metadata attn_metadata = metadata_cls( - seq_lens=torch.tensor(sequence_lengths, dtype=torch.int), - num_contexts=len(context_sequence_lengths), + seq_lens=torch.tensor([input_ids.size(-1)], dtype=torch.int), + num_contexts=1, kv_cache_params=KVCacheParams( use_cache=True, - num_cached_tokens_per_seq=past_seen_tokens, + num_cached_tokens_per_seq=num_cached_tokens_per_seq, ), + max_num_requests=1, + max_num_tokens=8192, kv_cache_manager=kv_cache_manager, request_ids=request_ids, prompt_lens=prompt_lens, - max_num_requests=len(context_sequence_lengths) + 2, - max_num_tokens=8192, ) - position_ids = [] - for i, tokens in enumerate(past_seen_tokens): - seq_len = context_sequence_lengths[i] if i < len( - context_sequence_lengths) else 1 - position_id = torch.arange(tokens, - tokens + seq_len, - device=input_ids.device) - position_ids.append(position_id) - - position_ids = torch.cat(position_ids).unsqueeze(0) - + # prefill + position_ids = [torch.arange(0, input_ids.size(-1))] + position_ids = torch.cat(position_ids).unsqueeze(0).cuda() with torch.inference_mode(): attn_metadata.prepare() logits = nemotron_h.forward(input_ids=input_ids, position_ids=position_ids, - attn_metadata=attn_metadata) + attn_metadata=attn_metadata, + return_context_logits=True) + + # compute logprobs from logits + prefill_logprobs = get_logprobs(input_ids[1:], logits) + + # reference logprobs from mcore + prefill_logprobs_ref = torch.tensor([ + -7.415980815887451, -0.36192911863327026, -2.8658294677734375, + -2.316344738006592 + ], + device=device) + + # reference logprobs from initial implementation + prefill_logprobs_ref_initial = torch.tensor([ + -7.4359540939331055, -0.37661877274513245, -2.8925108909606934, + -2.268364906311035 + ], + device=device) + + # compare logprobs with mcore logprobs, check that the max error is less than 0.3 + torch.testing.assert_close(prefill_logprobs, + prefill_logprobs_ref, + atol=0.3, + rtol=0.0) + + # compare logprobs with initial implementation, check that the max error is less than 0.1 + torch.testing.assert_close(prefill_logprobs, + prefill_logprobs_ref_initial, + atol=0.1, + rtol=0.0) + + print("#" * 40 + " prefill logprobs error " + "#" * 40) + print(" prefill_logprobs:", prefill_logprobs.cpu().numpy()) + print(" prefill_logprobs_ref:", + prefill_logprobs_ref.cpu().numpy()) + print("prefill_logprobs_ref_initial:", + prefill_logprobs_ref_initial.cpu().numpy()) + print( + f" max error mcore: {torch.max(torch.abs(prefill_logprobs - prefill_logprobs_ref)).item():.5f}" + ) + print( + f"max error initial: {torch.max(torch.abs(prefill_logprobs - prefill_logprobs_ref_initial)).item():.5f}" + ) + print() + + # output tokens + output = [] + decode_logprobs = [] + + # sample token greedily + sampled_tokens = torch.argmax(logits[-1]).unsqueeze(0) + output.append(sampled_tokens) + + decode_logprobs.append(get_logprobs(sampled_tokens, logits[-1:])) + + # generate 8 more tokens + max_tokens = 8 + for i in range(max_tokens): + + num_cached_tokens_per_seq = input_ids.shape[0] + i + 1 - self.assertEqual(len(past_seen_tokens), logits.shape[0]) + attn_metadata = metadata_cls( + seq_lens=torch.tensor([1], dtype=torch.int), + num_contexts=0, + kv_cache_params=KVCacheParams( + use_cache=True, + num_cached_tokens_per_seq=num_cached_tokens_per_seq, + ), + max_num_requests=1, + max_num_tokens=8192, + kv_cache_manager=kv_cache_manager, + request_ids=request_ids, + prompt_lens=prompt_lens, + ) + with torch.inference_mode(): + attn_metadata.prepare() + logits = nemotron_h.forward(input_ids=sampled_tokens, + position_ids=position_ids, + attn_metadata=attn_metadata) + + sampled_tokens = torch.argmax(logits).unsqueeze(0) + output.append(sampled_tokens) + + decode_logprobs.append(get_logprobs(sampled_tokens, logits)) + + output = torch.cat(output) + decode_logprobs = torch.cat(decode_logprobs) + completion = tokenizer.decode(output) + + decode_logprobs_ref_initial = torch.tensor([ + -2.2722280025482178, -0.5235245823860168, -0.8821321725845337, + -1.9436249732971191, -0.07366813719272614, -0.4224405586719513, + -0.3872227966785431, -0.0612114779651165, -1.0475994348526 + ], + device=device) + + self.assertEqual( + output.tolist(), + [18168, 1044, 1454, 58096, 32975, 1394, 32492, 1321, 6762]) + self.assertEqual( + completion, + " bright, with endless possibilities for innovation and growth") + + # compare logprobs with initial implementation, check that the max error is less than 0.1 + torch.testing.assert_close(decode_logprobs, + decode_logprobs_ref_initial, + atol=0.1, + rtol=0.0) + + print("#" * 40 + " completion " + "#" * 40) + print("completion ids:", output.tolist()) + print("sequence:", f"The future of AI is{completion}") + print(" decode_logprobs:", decode_logprobs.cpu().numpy()) + print("decode_logprobs_ref_initial:", + decode_logprobs_ref_initial.cpu().numpy()) + print( + f"max decode error initial: {torch.max(torch.abs(decode_logprobs - decode_logprobs_ref_initial)).item():.5f}" + ) + print() + + # now let's test that decodes match prefill logprobs + + input_ids = torch.cat([input_ids, output]) + prefill_decode_logprobs = torch.cat([prefill_logprobs, decode_logprobs]) + num_cached_tokens_per_seq = [0] + request_ids = [1] + prompt_lens = [input_ids.size(-1)] + + attn_metadata = metadata_cls( + seq_lens=torch.tensor([input_ids.size(-1)], dtype=torch.int), + num_contexts=1, + kv_cache_params=KVCacheParams( + use_cache=True, + num_cached_tokens_per_seq=num_cached_tokens_per_seq, + ), + max_num_requests=1, + max_num_tokens=8192, + kv_cache_manager=kv_cache_manager, + request_ids=request_ids, + prompt_lens=prompt_lens, + ) + + # prefill + position_ids = [torch.arange(0, input_ids.size(-1))] + position_ids = torch.cat(position_ids).unsqueeze(0).cuda() with torch.inference_mode(): attn_metadata.prepare() logits = nemotron_h.forward(input_ids=input_ids, position_ids=position_ids, attn_metadata=attn_metadata, return_context_logits=True) - self.assertEqual(input_ids.shape, logits.shape[:-1]) + + # compute logprobs from logits + full_sequence_prefill_logprobs = get_logprobs(input_ids[1:], logits) + + # compare full sequence prefill logprobs with prefill + decode logprobs, check that the max error is less than 0.3 + torch.testing.assert_close(full_sequence_prefill_logprobs, + prefill_decode_logprobs, + atol=0.3, + rtol=0.0) + + print("#" * 40 + " prefill vs decode logprobs " + "#" * 40) + print("full_sequence_prefill_logprobs:", + full_sequence_prefill_logprobs.cpu().numpy()) + print(" prefill_decode_logprobs:", + prefill_decode_logprobs.cpu().numpy()) + print( + f"max error: {torch.max(torch.abs(full_sequence_prefill_logprobs - prefill_decode_logprobs)).item():.5f}" + ) kv_cache_manager.shutdown() From 93ecd2adb13a09fb43b487fe1549410af87d4eff Mon Sep 17 00:00:00 2001 From: Omer Ullman Argov <118735753+omera-nv@users.noreply.github.com> Date: Thu, 8 May 2025 07:22:20 +0000 Subject: [PATCH 2/9] skip if not enough memory Signed-off-by: Omer Ullman Argov <118735753+omera-nv@users.noreply.github.com> --- .../modeling/test_modeling_nemotron_h.py | 60 ++----------------- 1 file changed, 6 insertions(+), 54 deletions(-) diff --git a/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py b/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py index eff77b338a88..b8c93cb5e282 100644 --- a/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py +++ b/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py @@ -1,5 +1,4 @@ import unittest -from copy import deepcopy import torch @@ -15,6 +14,7 @@ # isort: on from transformers import AutoTokenizer from utils.llm_data import llm_models_root +from utils.util import skip_gpu_memory_less_than from tensorrt_llm._torch.pyexecutor.model_engine import load_weights from tensorrt_llm._torch.pyexecutor.resource_manager import \ @@ -22,51 +22,6 @@ from tensorrt_llm.bindings.executor import KvCacheConfig from tensorrt_llm.mapping import Mapping -NEMOTRON_H_CONFIG = { - "architectures": ["NemotronHForCausalLM"], - "attention_bias": False, - "attention_dropout": 0.0, - "attention_head_dim": 128, - "bos_token_id": 1, - "chunk_size": 256, - "conv_kernel": 4, - "eos_token_id": 2, - "expand": 2, - "hidden_dropout": 0.0, - "hidden_size": 4096, - "hybrid_override_pattern": - "M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M-", - "initializer_range": 0.02, - "intermediate_size": 21504, - "layer_norm_epsilon": 1e-05, - "mamba_head_dim": 64, - "mamba_hidden_act": "silu", - "mamba_num_heads": 128, - "mamba_proj_bias": False, - "max_position_embeddings": 8192, - "mlp_bias": False, - "mlp_hidden_act": "relu2", - "model_type": "nemotron_h", - "n_groups": 8, - "num_attention_heads": 32, - "num_hidden_layers": 52, - "num_key_value_heads": 8, - "num_logits_to_keep": 1, - "pad_token_id": 0, - "rescale_prenorm_residual": True, - "residual_in_fp32": False, - "rms_norm_eps": 1e-05, - "sliding_window": None, - "ssm_state_size": 128, - "tie_word_embeddings": False, - "torch_dtype": "bfloat16", - "use_bias": False, - "use_cache": True, - "use_conv_bias": True, - "use_mamba_kernels": True, - "vocab_size": 131072 -} - def get_logprobs(token_ids: torch.Tensor, logits: torch.Tensor) -> torch.Tensor: raw_probs = torch.softmax(logits, dim=-1) @@ -77,16 +32,18 @@ def get_logprobs(token_ids: torch.Tensor, logits: torch.Tensor) -> torch.Tensor: class TestNemotronH(unittest.TestCase): + @skip_gpu_memory_less_than( + (2 * 8 + 1) * 2**30) # 8B, bf16, plus 1 GB for good measure def test_nemotron_correctness(self): - config_dict = deepcopy(NEMOTRON_H_CONFIG) - nemotron_h_config = NemotronHConfig.from_dict(config_dict) - model_dir = f"{llm_models_root(check=True)}/Nemotron-H-8B-Base-8K" + nemotron_h_config = NemotronHConfig.from_pretrained(model_dir) tokenizer = AutoTokenizer.from_pretrained(model_dir) dtype = nemotron_h_config.torch_dtype device = torch.device('cuda') + assert dtype == torch.bfloat16 + kv_cache_dtype = tensorrt_llm.bindings.DataType.BF16 model_config = ModelConfig(pretrained_config=nemotron_h_config) nemotron_h = NemotronHForCausalLM(model_config).to(device) @@ -113,11 +70,6 @@ def test_nemotron_correctness(self): max_seq_len = num_blocks * tokens_per_block max_batch_size = 1 - if dtype == torch.bfloat16: - kv_cache_dtype = tensorrt_llm.bindings.DataType.BF16 - else: - raise ValueError("Invalid dtype") - mapping = Mapping(world_size=1, tp_size=1, rank=0) kv_cache_config = KvCacheConfig(max_tokens=num_blocks * tokens_per_block, From b5296053bad7813888ec1fecb28859f8eb57f436 Mon Sep 17 00:00:00 2001 From: Omer Ullman Argov <118735753+omera-nv@users.noreply.github.com> Date: Mon, 12 May 2025 08:32:42 +0000 Subject: [PATCH 3/9] refactor test and remove prints Signed-off-by: Omer Ullman Argov <118735753+omera-nv@users.noreply.github.com> --- .../modeling/test_modeling_nemotron_h.py | 298 ++++++------------ 1 file changed, 105 insertions(+), 193 deletions(-) diff --git a/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py b/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py index b8c93cb5e282..f5c5af7380c1 100644 --- a/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py +++ b/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py @@ -30,6 +30,96 @@ def get_logprobs(token_ids: torch.Tensor, logits: torch.Tensor) -> torch.Tensor: return torch.log(token_probs) +def generate(model: NemotronHForCausalLM, cache: MambaHybridCacheManager, + prompt: list[int], tokens_to_generate: int, + device: torch.device) -> tuple[list[int], list[float]]: + """ + Generate `tokens_to_generate` tokens from the given prompt using the given model and cache. + Return the prefill+generated tokens and their logprobs, minus the first token in the prompt. + """ + request_ids = torch.randint(1, 100, (1, )).tolist() + token_nums = [len(prompt)] + num_cached_tokens = 0 + input_ids = torch.tensor(prompt, dtype=torch.int64, device=device) + position_ids = [torch.arange(0, input_ids.size(-1))] + tokens = prompt.copy() + + requests = cache.add_dummy_requests(request_ids, token_nums) + cache.prepare_mamba_cache_blocks(request_ids) + + metadata_cls = get_attention_backend( + model.model_config.attn_backend).Metadata + attn_metadata = metadata_cls( + seq_lens=torch.tensor([input_ids.size(-1)], dtype=torch.int), + num_contexts=1, + kv_cache_params=KVCacheParams( + use_cache=True, + num_cached_tokens_per_seq=[num_cached_tokens], + ), + max_num_requests=1, + max_num_tokens=8192, + kv_cache_manager=cache, + request_ids=request_ids, + prompt_lens=token_nums, + ) + + # prefill + position_ids = torch.cat(position_ids).unsqueeze(0).cuda() + with torch.inference_mode(): + attn_metadata.prepare() + logits = model.forward(input_ids=input_ids, + position_ids=position_ids, + attn_metadata=attn_metadata, + return_context_logits=True) + + # compute logprobs from logits + prefill_logprobs = get_logprobs(input_ids[1:], logits[:-1]) + logprobs = prefill_logprobs.tolist() + + # sample token greedily + sampled_tokens = torch.argmax(logits[-1]).unsqueeze(0) + tokens.append(sampled_tokens.item()) + logprobs.append(get_logprobs(sampled_tokens, logits[-1:]).item()) + + # one token already generated at prefill + for _ in range(tokens_to_generate - 1): + num_cached_tokens += input_ids.shape[0] + input_ids = sampled_tokens + position_ids = torch.tensor([num_cached_tokens], + dtype=torch.int64, + device=device) + + attn_metadata = metadata_cls( + seq_lens=torch.tensor([1], dtype=torch.int), + num_contexts=0, + kv_cache_params=KVCacheParams( + use_cache=True, + num_cached_tokens_per_seq=[num_cached_tokens], + ), + max_num_requests=1, + max_num_tokens=8192, + kv_cache_manager=cache, + request_ids=request_ids, + prompt_lens=token_nums, + ) + + with torch.inference_mode(): + attn_metadata.prepare() + logits = model.forward(input_ids=input_ids, + position_ids=position_ids, + attn_metadata=attn_metadata) + + # sample token greedily + sampled_tokens = torch.argmax(logits[-1]).unsqueeze(0) + tokens.append(sampled_tokens.item()) + logprobs.append(get_logprobs(sampled_tokens, logits[-1:]).item()) + + for req in requests: + cache.free_resources(req) + + return tokens, logprobs + + class TestNemotronH(unittest.TestCase): @skip_gpu_memory_less_than( @@ -51,15 +141,8 @@ def test_nemotron_correctness(self): weights = load_weights(model_dir) nemotron_h.load_weights(weights) - # These tokens are from the sentence: The future of AI is - input_ids = torch.tensor([1784, 7147, 1307, 26554, 1395], - dtype=torch.int64, - device=device) - - num_cached_tokens_per_seq = [0] - request_ids = [1] - token_nums = [input_ids.size(-1)] - prompt_lens = [input_ids.size(-1)] + text_prompt = "The future of AI is" + prompt = tokenizer.encode(text_prompt, add_special_tokens=False) num_blocks = 100 tokens_per_block = 128 @@ -97,201 +180,30 @@ def test_nemotron_correctness(self): dtype=kv_cache_dtype, num_extra_kv_tokens=0, ) - kv_cache_manager.add_dummy_requests(request_ids, token_nums) - kv_cache_manager.prepare_mamba_cache_blocks(request_ids) - metadata_cls = get_attention_backend(model_config.attn_backend).Metadata - attn_metadata = metadata_cls( - seq_lens=torch.tensor([input_ids.size(-1)], dtype=torch.int), - num_contexts=1, - kv_cache_params=KVCacheParams( - use_cache=True, - num_cached_tokens_per_seq=num_cached_tokens_per_seq, - ), - max_num_requests=1, - max_num_tokens=8192, - kv_cache_manager=kv_cache_manager, - request_ids=request_ids, - prompt_lens=prompt_lens, - ) - - # prefill - position_ids = [torch.arange(0, input_ids.size(-1))] - position_ids = torch.cat(position_ids).unsqueeze(0).cuda() - with torch.inference_mode(): - attn_metadata.prepare() - logits = nemotron_h.forward(input_ids=input_ids, - position_ids=position_ids, - attn_metadata=attn_metadata, - return_context_logits=True) - - # compute logprobs from logits - prefill_logprobs = get_logprobs(input_ids[1:], logits) + tokens, logprobs = generate(model=nemotron_h, + cache=kv_cache_manager, + prompt=prompt, + tokens_to_generate=9, + device=torch.device("cuda")) - # reference logprobs from mcore + # reference logprobs from mcore for prompt minus first token + # TODO(oargov): generate a reference on-the-fly once we have confidence in the HF impl prefill_logprobs_ref = torch.tensor([ -7.415980815887451, -0.36192911863327026, -2.8658294677734375, -2.316344738006592 - ], - device=device) - - # reference logprobs from initial implementation - prefill_logprobs_ref_initial = torch.tensor([ - -7.4359540939331055, -0.37661877274513245, -2.8925108909606934, - -2.268364906311035 - ], - device=device) + ]) # compare logprobs with mcore logprobs, check that the max error is less than 0.3 - torch.testing.assert_close(prefill_logprobs, + torch.testing.assert_close(torch.tensor(logprobs[:len(prompt) - 1]), prefill_logprobs_ref, atol=0.3, rtol=0.0) - # compare logprobs with initial implementation, check that the max error is less than 0.1 - torch.testing.assert_close(prefill_logprobs, - prefill_logprobs_ref_initial, - atol=0.1, - rtol=0.0) - - print("#" * 40 + " prefill logprobs error " + "#" * 40) - print(" prefill_logprobs:", prefill_logprobs.cpu().numpy()) - print(" prefill_logprobs_ref:", - prefill_logprobs_ref.cpu().numpy()) - print("prefill_logprobs_ref_initial:", - prefill_logprobs_ref_initial.cpu().numpy()) - print( - f" max error mcore: {torch.max(torch.abs(prefill_logprobs - prefill_logprobs_ref)).item():.5f}" - ) - print( - f"max error initial: {torch.max(torch.abs(prefill_logprobs - prefill_logprobs_ref_initial)).item():.5f}" - ) - print() - - # output tokens - output = [] - decode_logprobs = [] - - # sample token greedily - sampled_tokens = torch.argmax(logits[-1]).unsqueeze(0) - output.append(sampled_tokens) - - decode_logprobs.append(get_logprobs(sampled_tokens, logits[-1:])) - - # generate 8 more tokens - max_tokens = 8 - for i in range(max_tokens): - - num_cached_tokens_per_seq = input_ids.shape[0] + i + 1 - - attn_metadata = metadata_cls( - seq_lens=torch.tensor([1], dtype=torch.int), - num_contexts=0, - kv_cache_params=KVCacheParams( - use_cache=True, - num_cached_tokens_per_seq=num_cached_tokens_per_seq, - ), - max_num_requests=1, - max_num_tokens=8192, - kv_cache_manager=kv_cache_manager, - request_ids=request_ids, - prompt_lens=prompt_lens, - ) - - with torch.inference_mode(): - attn_metadata.prepare() - logits = nemotron_h.forward(input_ids=sampled_tokens, - position_ids=position_ids, - attn_metadata=attn_metadata) - - sampled_tokens = torch.argmax(logits).unsqueeze(0) - output.append(sampled_tokens) - - decode_logprobs.append(get_logprobs(sampled_tokens, logits)) - - output = torch.cat(output) - decode_logprobs = torch.cat(decode_logprobs) - completion = tokenizer.decode(output) - - decode_logprobs_ref_initial = torch.tensor([ - -2.2722280025482178, -0.5235245823860168, -0.8821321725845337, - -1.9436249732971191, -0.07366813719272614, -0.4224405586719513, - -0.3872227966785431, -0.0612114779651165, -1.0475994348526 - ], - device=device) - - self.assertEqual( - output.tolist(), - [18168, 1044, 1454, 58096, 32975, 1394, 32492, 1321, 6762]) + # TODO(oargov): get a reference for the decode logprobs and compare + output = tokenizer.decode(tokens[len(prompt) - 1:]) self.assertEqual( - completion, - " bright, with endless possibilities for innovation and growth") - - # compare logprobs with initial implementation, check that the max error is less than 0.1 - torch.testing.assert_close(decode_logprobs, - decode_logprobs_ref_initial, - atol=0.1, - rtol=0.0) - - print("#" * 40 + " completion " + "#" * 40) - print("completion ids:", output.tolist()) - print("sequence:", f"The future of AI is{completion}") - print(" decode_logprobs:", decode_logprobs.cpu().numpy()) - print("decode_logprobs_ref_initial:", - decode_logprobs_ref_initial.cpu().numpy()) - print( - f"max decode error initial: {torch.max(torch.abs(decode_logprobs - decode_logprobs_ref_initial)).item():.5f}" - ) - print() - - # now let's test that decodes match prefill logprobs - - input_ids = torch.cat([input_ids, output]) - prefill_decode_logprobs = torch.cat([prefill_logprobs, decode_logprobs]) - num_cached_tokens_per_seq = [0] - request_ids = [1] - prompt_lens = [input_ids.size(-1)] - - attn_metadata = metadata_cls( - seq_lens=torch.tensor([input_ids.size(-1)], dtype=torch.int), - num_contexts=1, - kv_cache_params=KVCacheParams( - use_cache=True, - num_cached_tokens_per_seq=num_cached_tokens_per_seq, - ), - max_num_requests=1, - max_num_tokens=8192, - kv_cache_manager=kv_cache_manager, - request_ids=request_ids, - prompt_lens=prompt_lens, - ) - - # prefill - position_ids = [torch.arange(0, input_ids.size(-1))] - position_ids = torch.cat(position_ids).unsqueeze(0).cuda() - with torch.inference_mode(): - attn_metadata.prepare() - logits = nemotron_h.forward(input_ids=input_ids, - position_ids=position_ids, - attn_metadata=attn_metadata, - return_context_logits=True) - - # compute logprobs from logits - full_sequence_prefill_logprobs = get_logprobs(input_ids[1:], logits) - - # compare full sequence prefill logprobs with prefill + decode logprobs, check that the max error is less than 0.3 - torch.testing.assert_close(full_sequence_prefill_logprobs, - prefill_decode_logprobs, - atol=0.3, - rtol=0.0) - - print("#" * 40 + " prefill vs decode logprobs " + "#" * 40) - print("full_sequence_prefill_logprobs:", - full_sequence_prefill_logprobs.cpu().numpy()) - print(" prefill_decode_logprobs:", - prefill_decode_logprobs.cpu().numpy()) - print( - f"max error: {torch.max(torch.abs(full_sequence_prefill_logprobs - prefill_decode_logprobs)).item():.5f}" - ) + output, + " is bright, with endless possibilities for innovation and growth") kv_cache_manager.shutdown() From 05de495038427db6888fc0456909b2a1a796f19a Mon Sep 17 00:00:00 2001 From: Omer Ullman Argov <118735753+omera-nv@users.noreply.github.com> Date: Thu, 15 May 2025 10:30:10 +0000 Subject: [PATCH 4/9] fix for new api Signed-off-by: Omer Ullman Argov <118735753+omera-nv@users.noreply.github.com> --- tests/unittest/_torch/modeling/test_modeling_nemotron_h.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py b/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py index f5c5af7380c1..ae09ec4b3e3e 100644 --- a/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py +++ b/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py @@ -138,7 +138,8 @@ def test_nemotron_correctness(self): model_config = ModelConfig(pretrained_config=nemotron_h_config) nemotron_h = NemotronHForCausalLM(model_config).to(device) - weights = load_weights(model_dir) + mapping = Mapping(world_size=1, tp_size=1, rank=0) + weights = load_weights(model_dir, mapping=mapping) nemotron_h.load_weights(weights) text_prompt = "The future of AI is" @@ -153,7 +154,6 @@ def test_nemotron_correctness(self): max_seq_len = num_blocks * tokens_per_block max_batch_size = 1 - mapping = Mapping(world_size=1, tp_size=1, rank=0) kv_cache_config = KvCacheConfig(max_tokens=num_blocks * tokens_per_block, enable_block_reuse=False) From 516b3cf0f21cfaabe14e7799bc5f27e9877a527d Mon Sep 17 00:00:00 2001 From: Tomer Asida <57313761+tomeras91@users.noreply.github.com> Date: Sun, 18 May 2025 11:30:38 +0000 Subject: [PATCH 5/9] adapt test to work for batched requests as well Signed-off-by: Tomer Asida <57313761+tomeras91@users.noreply.github.com> --- .../modeling/test_modeling_nemotron_h.py | 133 +++++++++++------- 1 file changed, 83 insertions(+), 50 deletions(-) diff --git a/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py b/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py index ae09ec4b3e3e..0c6c776dfd8e 100644 --- a/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py +++ b/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py @@ -12,7 +12,7 @@ NemotronHForCausalLM ) # isort: on -from transformers import AutoTokenizer +from transformers import AutoTokenizer, PreTrainedTokenizerBase from utils.llm_data import llm_models_root from utils.util import skip_gpu_memory_less_than @@ -30,40 +30,49 @@ def get_logprobs(token_ids: torch.Tensor, logits: torch.Tensor) -> torch.Tensor: return torch.log(token_probs) -def generate(model: NemotronHForCausalLM, cache: MambaHybridCacheManager, - prompt: list[int], tokens_to_generate: int, - device: torch.device) -> tuple[list[int], list[float]]: +def generate( + model: NemotronHForCausalLM, tokenizer: PreTrainedTokenizerBase, + cache: MambaHybridCacheManager, text_prompts: list[str], + tokens_to_generate: int, device: torch.device +) -> tuple[list[int], list[list[int]], list[list[float]]]: """ - Generate `tokens_to_generate` tokens from the given prompt using the given model and cache. - Return the prefill+generated tokens and their logprobs, minus the first token in the prompt. + Generate `tokens_to_generate` tokens from the given prompts using the given model and cache. + Return the prompt_lens along with the prefill+generated tokens and their logprobs, minus the first token in the prompt. """ - request_ids = torch.randint(1, 100, (1, )).tolist() - token_nums = [len(prompt)] - num_cached_tokens = 0 - input_ids = torch.tensor(prompt, dtype=torch.int64, device=device) - position_ids = [torch.arange(0, input_ids.size(-1))] - tokens = prompt.copy() - - requests = cache.add_dummy_requests(request_ids, token_nums) + num_seqs = len(text_prompts) + all_token_ids = [ + tokenizer.encode(prompt, add_special_tokens=False) + for prompt in text_prompts + ] + input_ids = torch.cat([ + torch.tensor(token_ids, dtype=torch.int64, device=device) + for token_ids in all_token_ids + ], + dim=0) + request_ids = list(range(1, num_seqs + 1)) + prompt_lens = [len(token_ids) for token_ids in all_token_ids] + + requests = cache.add_dummy_requests(request_ids, prompt_lens) cache.prepare_mamba_cache_blocks(request_ids) metadata_cls = get_attention_backend( model.model_config.attn_backend).Metadata attn_metadata = metadata_cls( - seq_lens=torch.tensor([input_ids.size(-1)], dtype=torch.int), - num_contexts=1, + seq_lens=torch.tensor(prompt_lens, dtype=torch.int), + num_contexts=num_seqs, kv_cache_params=KVCacheParams( use_cache=True, - num_cached_tokens_per_seq=[num_cached_tokens], + num_cached_tokens_per_seq=[0] * num_seqs, ), - max_num_requests=1, + max_num_requests=num_seqs, max_num_tokens=8192, kv_cache_manager=cache, request_ids=request_ids, - prompt_lens=token_nums, + prompt_lens=prompt_lens, ) # prefill + position_ids = [torch.arange(0, prompt_len) for prompt_len in prompt_lens] position_ids = torch.cat(position_ids).unsqueeze(0).cuda() with torch.inference_mode(): attn_metadata.prepare() @@ -73,51 +82,66 @@ def generate(model: NemotronHForCausalLM, cache: MambaHybridCacheManager, return_context_logits=True) # compute logprobs from logits - prefill_logprobs = get_logprobs(input_ids[1:], logits[:-1]) - logprobs = prefill_logprobs.tolist() + all_logits = logits.split(prompt_lens, dim=0) + all_logprobs = [ + get_logprobs( + torch.tensor(token_ids[1:], dtype=torch.int64, device=device), + this_logits[:-1]).tolist() + for token_ids, this_logits in zip(all_token_ids, all_logits) + ] # sample token greedily - sampled_tokens = torch.argmax(logits[-1]).unsqueeze(0) - tokens.append(sampled_tokens.item()) - logprobs.append(get_logprobs(sampled_tokens, logits[-1:]).item()) + sampled_tokens = torch.cat([ + torch.argmax(this_logits[-1]).unsqueeze(0) for this_logits in all_logits + ], + dim=0) + for i in range(num_seqs): + all_token_ids[i].append(sampled_tokens[i].item()) + all_logprobs[i].append( + get_logprobs(sampled_tokens[i].unsqueeze(0), + all_logits[i][-1:]).item()) # one token already generated at prefill - for _ in range(tokens_to_generate - 1): - num_cached_tokens += input_ids.shape[0] - input_ids = sampled_tokens - position_ids = torch.tensor([num_cached_tokens], + for i in range(tokens_to_generate - 1): + num_cached_tokens_per_seq = [ + prompt_len + i + 1 for prompt_len in prompt_lens + ] + position_ids = torch.tensor([num_cached_tokens_per_seq], dtype=torch.int64, device=device) attn_metadata = metadata_cls( - seq_lens=torch.tensor([1], dtype=torch.int), + seq_lens=torch.tensor([1] * num_seqs, dtype=torch.int), num_contexts=0, kv_cache_params=KVCacheParams( use_cache=True, - num_cached_tokens_per_seq=[num_cached_tokens], + num_cached_tokens_per_seq=num_cached_tokens_per_seq, ), - max_num_requests=1, + max_num_requests=num_seqs, max_num_tokens=8192, kv_cache_manager=cache, request_ids=request_ids, - prompt_lens=token_nums, + prompt_lens=prompt_lens, ) with torch.inference_mode(): attn_metadata.prepare() - logits = model.forward(input_ids=input_ids, + logits = model.forward(input_ids=sampled_tokens, position_ids=position_ids, attn_metadata=attn_metadata) # sample token greedily - sampled_tokens = torch.argmax(logits[-1]).unsqueeze(0) - tokens.append(sampled_tokens.item()) - logprobs.append(get_logprobs(sampled_tokens, logits[-1:]).item()) + sampled_tokens = torch.argmax(logits, dim=-1, keepdim=False) + for i in range(num_seqs): + all_token_ids[i].append(sampled_tokens[i].item()) + all_logprobs[i].append( + get_logprobs(sampled_tokens[i].unsqueeze(0), + logits[i].unsqueeze(0)).item()) for req in requests: cache.free_resources(req) - return tokens, logprobs + return prompt_lens, all_token_ids, all_logprobs class TestNemotronH(unittest.TestCase): @@ -142,8 +166,10 @@ def test_nemotron_correctness(self): weights = load_weights(model_dir, mapping=mapping) nemotron_h.load_weights(weights) - text_prompt = "The future of AI is" - prompt = tokenizer.encode(text_prompt, add_special_tokens=False) + text_prompts = [ + "The future of AI is", + ] + num_prompts = len(text_prompts) num_blocks = 100 tokens_per_block = 128 @@ -152,7 +178,7 @@ def test_nemotron_correctness(self): mamba_num_layers = nemotron_h.config.hybrid_override_pattern.count("M") num_kv_heads = nemotron_h.config.num_key_value_heads max_seq_len = num_blocks * tokens_per_block - max_batch_size = 1 + max_batch_size = num_prompts kv_cache_config = KvCacheConfig(max_tokens=num_blocks * tokens_per_block, @@ -181,11 +207,12 @@ def test_nemotron_correctness(self): num_extra_kv_tokens=0, ) - tokens, logprobs = generate(model=nemotron_h, - cache=kv_cache_manager, - prompt=prompt, - tokens_to_generate=9, - device=torch.device("cuda")) + prompt_lens, tokens, logprobs = generate(model=nemotron_h, + tokenizer=tokenizer, + cache=kv_cache_manager, + text_prompts=text_prompts, + tokens_to_generate=9, + device=torch.device("cuda")) # reference logprobs from mcore for prompt minus first token # TODO(oargov): generate a reference on-the-fly once we have confidence in the HF impl @@ -195,15 +222,21 @@ def test_nemotron_correctness(self): ]) # compare logprobs with mcore logprobs, check that the max error is less than 0.3 - torch.testing.assert_close(torch.tensor(logprobs[:len(prompt) - 1]), + torch.testing.assert_close(torch.tensor(logprobs[0][:prompt_lens[0] - + 1]), prefill_logprobs_ref, atol=0.3, rtol=0.0) # TODO(oargov): get a reference for the decode logprobs and compare - output = tokenizer.decode(tokens[len(prompt) - 1:]) - self.assertEqual( - output, - " is bright, with endless possibilities for innovation and growth") + expected_completions = [ + " bright, with endless possibilities for innovation and growth", + ] + completions = [ + tokenizer.decode(tokens[i][prompt_lens[i]:]) + for i in range(num_prompts) + ] + for i in range(num_prompts): + self.assertEqual(completions[i], expected_completions[i]) kv_cache_manager.shutdown() From e5b2f88d4181b40008ce411af0cfbf62add07b9b Mon Sep 17 00:00:00 2001 From: Tomer Asida <57313761+tomeras91@users.noreply.github.com> Date: Sun, 18 May 2025 13:33:43 +0000 Subject: [PATCH 6/9] Add prefill+decode reference logprobs from initial implementation + batched forward test Signed-off-by: Tomer Asida <57313761+tomeras91@users.noreply.github.com> --- .../modeling/test_modeling_nemotron_h.py | 194 ++++++++++++++++-- 1 file changed, 172 insertions(+), 22 deletions(-) diff --git a/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py b/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py index 0c6c776dfd8e..e7b416aaf755 100644 --- a/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py +++ b/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py @@ -30,15 +30,11 @@ def get_logprobs(token_ids: torch.Tensor, logits: torch.Tensor) -> torch.Tensor: return torch.log(token_probs) -def generate( +def _generate( model: NemotronHForCausalLM, tokenizer: PreTrainedTokenizerBase, cache: MambaHybridCacheManager, text_prompts: list[str], tokens_to_generate: int, device: torch.device ) -> tuple[list[int], list[list[int]], list[list[float]]]: - """ - Generate `tokens_to_generate` tokens from the given prompts using the given model and cache. - Return the prompt_lens along with the prefill+generated tokens and their logprobs, minus the first token in the prompt. - """ num_seqs = len(text_prompts) all_token_ids = [ tokenizer.encode(prompt, add_special_tokens=False) @@ -144,6 +140,33 @@ def generate( return prompt_lens, all_token_ids, all_logprobs +def generate( + model: NemotronHForCausalLM, + tokenizer: PreTrainedTokenizerBase, + cache: MambaHybridCacheManager, + text_prompts: list[str], + tokens_to_generate: int, + device: torch.device, + one_by_one: bool = False +) -> tuple[list[int], list[list[int]], list[list[float]]]: + """ + Generate `tokens_to_generate` tokens from the given prompts using the given model and cache. + Return the prompt_lens along with the prefill+generated tokens and their logprobs, minus the first token in the prompt. + """ + if one_by_one: + num_prompts = len(text_prompts) + prompt_lens, tokens, logprobs = [None] * num_prompts, [ + None + ] * num_prompts, [None] * num_prompts + for i in range(num_prompts): + p, t, l = _generate(model, tokenizer, cache, [text_prompts[i]], + tokens_to_generate, device) + prompt_lens[i], tokens[i], logprobs[i] = p[0], t[0], l[0] + return prompt_lens, tokens, logprobs + return _generate(model, tokenizer, cache, text_prompts, tokens_to_generate, + device) + + class TestNemotronH(unittest.TestCase): @skip_gpu_memory_less_than( @@ -168,6 +191,7 @@ def test_nemotron_correctness(self): text_prompts = [ "The future of AI is", + "The president of the United States is", ] num_prompts = len(text_prompts) @@ -207,36 +231,162 @@ def test_nemotron_correctness(self): num_extra_kv_tokens=0, ) - prompt_lens, tokens, logprobs = generate(model=nemotron_h, - tokenizer=tokenizer, - cache=kv_cache_manager, - text_prompts=text_prompts, - tokens_to_generate=9, - device=torch.device("cuda")) + prompt_lens, tokens_no_batching, logprobs_no_batching = generate( + model=nemotron_h, + tokenizer=tokenizer, + cache=kv_cache_manager, + text_prompts=text_prompts, + tokens_to_generate=9, + device=torch.device("cuda"), + one_by_one=True) + completions_no_batching = [ + tokenizer.decode(tokens_no_batching[i][prompt_lens[i]:]) + for i in range(num_prompts) + ] - # reference logprobs from mcore for prompt minus first token + _, tokens_batching, logprobs_batching = generate( + model=nemotron_h, + tokenizer=tokenizer, + cache=kv_cache_manager, + text_prompts=text_prompts, + tokens_to_generate=9, + device=torch.device("cuda")) + completions_batching = [ + tokenizer.decode(tokens_batching[i][prompt_lens[i]:]) + for i in range(num_prompts) + ] + + # reference logprobs for first prompt from mcore for prompt minus first token # TODO(oargov): generate a reference on-the-fly once we have confidence in the HF impl - prefill_logprobs_ref = torch.tensor([ + prefill_logprobs_ref_mcore = torch.tensor([ -7.415980815887451, -0.36192911863327026, -2.8658294677734375, -2.316344738006592 ]) # compare logprobs with mcore logprobs, check that the max error is less than 0.3 - torch.testing.assert_close(torch.tensor(logprobs[0][:prompt_lens[0] - - 1]), - prefill_logprobs_ref, + torch.testing.assert_close(torch.tensor( + logprobs_no_batching[0][:prompt_lens[0] - 1]), + prefill_logprobs_ref_mcore, atol=0.3, rtol=0.0) - # TODO(oargov): get a reference for the decode logprobs and compare + # reference logprobs for first prompt from initial implementation (commit 5ce1102a02bd2938c0c8334138371f081f55fcc1) + prefill_logprobs_ref_initial_no_batching = [ + torch.tensor([ + -7.4359540939331055, + -0.37661877274513245, + -2.8925108909606934, + -2.268364906311035, + ]), + torch.tensor([ + -8.759482383728027, + -1.656238079071045, + -0.5448741912841797, + -1.7702054977416992, + -0.05832016468048096, + -1.460732102394104, + ]) + ] + prefill_logprobs_ref_initial_with_batching = [ + torch.tensor([ + -7.401950836181641, -0.38696032762527466, -2.8725428581237793, + -2.2654521465301514 + ]), + torch.tensor([ + -8.73007583618164, -1.6853574514389038, -0.5468529462814331, + -1.7846013307571411, -0.053610533475875854, -1.4385275840759277 + ]) + ] + prefill_batching_atol = min( + 0.1, 1.5 * max( + torch.max(torch.abs(x - y)) + for x, y in zip(prefill_logprobs_ref_initial_with_batching, + prefill_logprobs_ref_initial_no_batching))) + + decode_logprobs_ref_initial_no_batching = [ + torch.tensor([ + -2.2722280025482178, -0.5235245823860168, -0.8821321725845337, + -1.9436249732971191, -0.07366813719272614, -0.4224405586719513, + -0.3872227966785431, -0.06121065467596054, -1.0475994348526 + ]), + torch.tensor([ + -1.329713225364685, -1.6879069805145264, -0.040034178644418716, + -0.4808207154273987, -0.3581068515777588, -0.2784178853034973, + -0.005814795847982168, -0.0563097707927227, -0.05941024422645569 + ]) + ] + decode_logprobs_ref_initial_with_batching = [ + torch.tensor([ + -2.2877156734466553, -0.507795512676239, -0.8313305377960205, + -1.940523386001587, -0.07369701564311981, -0.4190545976161957, + -0.4250463843345642, -0.061063338071107864, -1.046282410621643 + ]), + torch.tensor([ + -1.3567769527435303, -1.7291667461395264, -0.04527968540787697, + -0.4836069345474243, -0.3971801698207855, -0.2481495887041092, + -0.005787517875432968, -0.056093256920576096, + -0.058267030864953995 + ]) + ] + decode_batching_atol = min( + 0.1, 1.5 * max( + torch.max(torch.abs(x - y)) + for x, y in zip(decode_logprobs_ref_initial_with_batching, + decode_logprobs_ref_initial_no_batching))) + expected_completions = [ " bright, with endless possibilities for innovation and growth", + " the head of state and head of government of", ] - completions = [ - tokenizer.decode(tokens[i][prompt_lens[i]:]) - for i in range(num_prompts) - ] + for i in range(num_prompts): - self.assertEqual(completions[i], expected_completions[i]) + prefill_logprobs_no_batching = torch.tensor( + logprobs_no_batching[i][:prompt_lens[i] - 1]) + decode_logprobs_no_batching = torch.tensor( + logprobs_no_batching[i][prompt_lens[i] - 1:]) + + prefill_logprobs_batching = torch.tensor( + logprobs_batching[i][:prompt_lens[i] - 1]) + decode_logprobs_batching = torch.tensor( + logprobs_batching[i][prompt_lens[i] - 1:]) + + # compare prompt logprobs with initial implementation + torch.testing.assert_close( + prefill_logprobs_no_batching, + prefill_logprobs_ref_initial_no_batching[i], + atol=0.1, + rtol=0.0) + torch.testing.assert_close( + prefill_logprobs_batching, + prefill_logprobs_ref_initial_with_batching[i], + atol=0.1, + rtol=0.0) + + # compare expected completion + self.assertEqual(completions_batching[i], expected_completions[i]) + self.assertEqual(completions_no_batching[i], + expected_completions[i]) + + # compare decode logprobs with initial implementation + torch.testing.assert_close( + decode_logprobs_no_batching, + decode_logprobs_ref_initial_no_batching[i], + atol=0.1, + rtol=0.0) + torch.testing.assert_close( + decode_logprobs_batching, + decode_logprobs_ref_initial_with_batching[i], + atol=0.1, + rtol=0.0) + + # compare logprobs with and without batching + torch.testing.assert_close(prefill_logprobs_batching, + prefill_logprobs_no_batching, + atol=prefill_batching_atol, + rtol=0.0) + torch.testing.assert_close(decode_logprobs_batching, + decode_logprobs_no_batching, + atol=decode_batching_atol, + rtol=0.0) kv_cache_manager.shutdown() From 1a88b87738b58158ef4d51a6eaaba24ddcf97440 Mon Sep 17 00:00:00 2001 From: Tomer Asida <57313761+tomeras91@users.noreply.github.com> Date: Sun, 18 May 2025 13:42:00 +0000 Subject: [PATCH 7/9] Add testing that decode matches prefill - compare decode vs all prefilling the decoded tokens Signed-off-by: Tomer Asida <57313761+tomeras91@users.noreply.github.com> --- .../modeling/test_modeling_nemotron_h.py | 108 +++++++++++------- 1 file changed, 65 insertions(+), 43 deletions(-) diff --git a/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py b/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py index e7b416aaf755..82c83b5eb231 100644 --- a/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py +++ b/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py @@ -86,53 +86,55 @@ def _generate( for token_ids, this_logits in zip(all_token_ids, all_logits) ] - # sample token greedily - sampled_tokens = torch.cat([ - torch.argmax(this_logits[-1]).unsqueeze(0) for this_logits in all_logits - ], - dim=0) - for i in range(num_seqs): - all_token_ids[i].append(sampled_tokens[i].item()) - all_logprobs[i].append( - get_logprobs(sampled_tokens[i].unsqueeze(0), - all_logits[i][-1:]).item()) - - # one token already generated at prefill - for i in range(tokens_to_generate - 1): - num_cached_tokens_per_seq = [ - prompt_len + i + 1 for prompt_len in prompt_lens - ] - position_ids = torch.tensor([num_cached_tokens_per_seq], - dtype=torch.int64, - device=device) - - attn_metadata = metadata_cls( - seq_lens=torch.tensor([1] * num_seqs, dtype=torch.int), - num_contexts=0, - kv_cache_params=KVCacheParams( - use_cache=True, - num_cached_tokens_per_seq=num_cached_tokens_per_seq, - ), - max_num_requests=num_seqs, - max_num_tokens=8192, - kv_cache_manager=cache, - request_ids=request_ids, - prompt_lens=prompt_lens, - ) - - with torch.inference_mode(): - attn_metadata.prepare() - logits = model.forward(input_ids=sampled_tokens, - position_ids=position_ids, - attn_metadata=attn_metadata) - + if tokens_to_generate > 0: # sample token greedily - sampled_tokens = torch.argmax(logits, dim=-1, keepdim=False) + sampled_tokens = torch.cat([ + torch.argmax(this_logits[-1]).unsqueeze(0) + for this_logits in all_logits + ], + dim=0) for i in range(num_seqs): all_token_ids[i].append(sampled_tokens[i].item()) all_logprobs[i].append( get_logprobs(sampled_tokens[i].unsqueeze(0), - logits[i].unsqueeze(0)).item()) + all_logits[i][-1:]).item()) + + # one token already generated at prefill + for i in range(tokens_to_generate - 1): + num_cached_tokens_per_seq = [ + prompt_len + i + 1 for prompt_len in prompt_lens + ] + position_ids = torch.tensor([num_cached_tokens_per_seq], + dtype=torch.int64, + device=device) + + attn_metadata = metadata_cls( + seq_lens=torch.tensor([1] * num_seqs, dtype=torch.int), + num_contexts=0, + kv_cache_params=KVCacheParams( + use_cache=True, + num_cached_tokens_per_seq=num_cached_tokens_per_seq, + ), + max_num_requests=num_seqs, + max_num_tokens=8192, + kv_cache_manager=cache, + request_ids=request_ids, + prompt_lens=prompt_lens, + ) + + with torch.inference_mode(): + attn_metadata.prepare() + logits = model.forward(input_ids=sampled_tokens, + position_ids=position_ids, + attn_metadata=attn_metadata) + + # sample token greedily + sampled_tokens = torch.argmax(logits, dim=-1, keepdim=False) + for i in range(num_seqs): + all_token_ids[i].append(sampled_tokens[i].item()) + all_logprobs[i].append( + get_logprobs(sampled_tokens[i].unsqueeze(0), + logits[i].unsqueeze(0)).item()) for req in requests: cache.free_resources(req) @@ -379,7 +381,7 @@ def test_nemotron_correctness(self): atol=0.1, rtol=0.0) - # compare logprobs with and without batching + # compare logprobs with and without batching, tolerace by diff in initial implementation torch.testing.assert_close(prefill_logprobs_batching, prefill_logprobs_no_batching, atol=prefill_batching_atol, @@ -389,4 +391,24 @@ def test_nemotron_correctness(self): atol=decode_batching_atol, rtol=0.0) + # now let's test that decodes match prefill logprobs + text_prompts_with_completions = [ + f"{text_prompts[i]}{completions_batching[i]}" + for i in range(num_prompts) + ] + _, _, full_sequence_logprobs = generate( + model=nemotron_h, + tokenizer=tokenizer, + cache=kv_cache_manager, + text_prompts=text_prompts_with_completions, + tokens_to_generate=0, + device=torch.device("cuda")) + + # compare full sequence logprobs with prefill+decode logprobs, tolerance like mcore tolerance + for i in range(num_prompts): + torch.testing.assert_close(torch.tensor(full_sequence_logprobs[i]), + torch.tensor(logprobs_batching[i]), + atol=0.3, + rtol=0.0) + kv_cache_manager.shutdown() From 5984afced987e7996d48e9939960c052651952a0 Mon Sep 17 00:00:00 2001 From: Tomer Asida <57313761+tomeras91@users.noreply.github.com> Date: Sun, 18 May 2025 16:14:16 +0000 Subject: [PATCH 8/9] relax tolerance Signed-off-by: Tomer Asida <57313761+tomeras91@users.noreply.github.com> --- .../modeling/test_modeling_nemotron_h.py | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py b/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py index 82c83b5eb231..b0eda73fb854 100644 --- a/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py +++ b/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py @@ -266,13 +266,16 @@ def test_nemotron_correctness(self): ]) # compare logprobs with mcore logprobs, check that the max error is less than 0.3 + mcore_atol = 0.3 torch.testing.assert_close(torch.tensor( logprobs_no_batching[0][:prompt_lens[0] - 1]), prefill_logprobs_ref_mcore, - atol=0.3, + atol=mcore_atol, rtol=0.0) # reference logprobs for first prompt from initial implementation (commit 5ce1102a02bd2938c0c8334138371f081f55fcc1) + initial_impl_atol = 0.2 + prefill_logprobs_ref_initial_no_batching = [ torch.tensor([ -7.4359540939331055, @@ -300,7 +303,7 @@ def test_nemotron_correctness(self): ]) ] prefill_batching_atol = min( - 0.1, 1.5 * max( + initial_impl_atol, 1.5 * max( torch.max(torch.abs(x - y)) for x, y in zip(prefill_logprobs_ref_initial_with_batching, prefill_logprobs_ref_initial_no_batching))) @@ -331,7 +334,7 @@ def test_nemotron_correctness(self): ]) ] decode_batching_atol = min( - 0.1, 1.5 * max( + initial_impl_atol, 1.5 * max( torch.max(torch.abs(x - y)) for x, y in zip(decode_logprobs_ref_initial_with_batching, decode_logprobs_ref_initial_no_batching))) @@ -356,12 +359,12 @@ def test_nemotron_correctness(self): torch.testing.assert_close( prefill_logprobs_no_batching, prefill_logprobs_ref_initial_no_batching[i], - atol=0.1, + atol=initial_impl_atol, rtol=0.0) torch.testing.assert_close( prefill_logprobs_batching, prefill_logprobs_ref_initial_with_batching[i], - atol=0.1, + atol=initial_impl_atol, rtol=0.0) # compare expected completion @@ -373,12 +376,12 @@ def test_nemotron_correctness(self): torch.testing.assert_close( decode_logprobs_no_batching, decode_logprobs_ref_initial_no_batching[i], - atol=0.1, + atol=initial_impl_atol, rtol=0.0) torch.testing.assert_close( decode_logprobs_batching, decode_logprobs_ref_initial_with_batching[i], - atol=0.1, + atol=initial_impl_atol, rtol=0.0) # compare logprobs with and without batching, tolerace by diff in initial implementation @@ -408,7 +411,7 @@ def test_nemotron_correctness(self): for i in range(num_prompts): torch.testing.assert_close(torch.tensor(full_sequence_logprobs[i]), torch.tensor(logprobs_batching[i]), - atol=0.3, + atol=mcore_atol, rtol=0.0) kv_cache_manager.shutdown() From b30e068804806de991630528bdd51b3cf26f853e Mon Sep 17 00:00:00 2001 From: Tomer Asida <57313761+tomeras91@users.noreply.github.com> Date: Sun, 18 May 2025 20:09:44 +0000 Subject: [PATCH 9/9] relax tolerance Signed-off-by: Tomer Asida <57313761+tomeras91@users.noreply.github.com> --- .../_torch/modeling/test_modeling_nemotron_h.py | 17 ++++------------- 1 file changed, 4 insertions(+), 13 deletions(-) diff --git a/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py b/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py index b0eda73fb854..54af9c26946a 100644 --- a/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py +++ b/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py @@ -273,8 +273,9 @@ def test_nemotron_correctness(self): atol=mcore_atol, rtol=0.0) - # reference logprobs for first prompt from initial implementation (commit 5ce1102a02bd2938c0c8334138371f081f55fcc1) + # reference logprobs from initial implementation (commit 5ce1102a02bd2938c0c8334138371f081f55fcc1) initial_impl_atol = 0.2 + batching_atol = 0.2 prefill_logprobs_ref_initial_no_batching = [ torch.tensor([ @@ -302,11 +303,6 @@ def test_nemotron_correctness(self): -1.7846013307571411, -0.053610533475875854, -1.4385275840759277 ]) ] - prefill_batching_atol = min( - initial_impl_atol, 1.5 * max( - torch.max(torch.abs(x - y)) - for x, y in zip(prefill_logprobs_ref_initial_with_batching, - prefill_logprobs_ref_initial_no_batching))) decode_logprobs_ref_initial_no_batching = [ torch.tensor([ @@ -333,11 +329,6 @@ def test_nemotron_correctness(self): -0.058267030864953995 ]) ] - decode_batching_atol = min( - initial_impl_atol, 1.5 * max( - torch.max(torch.abs(x - y)) - for x, y in zip(decode_logprobs_ref_initial_with_batching, - decode_logprobs_ref_initial_no_batching))) expected_completions = [ " bright, with endless possibilities for innovation and growth", @@ -387,11 +378,11 @@ def test_nemotron_correctness(self): # compare logprobs with and without batching, tolerace by diff in initial implementation torch.testing.assert_close(prefill_logprobs_batching, prefill_logprobs_no_batching, - atol=prefill_batching_atol, + atol=batching_atol, rtol=0.0) torch.testing.assert_close(decode_logprobs_batching, decode_logprobs_no_batching, - atol=decode_batching_atol, + atol=batching_atol, rtol=0.0) # now let's test that decodes match prefill logprobs