Skip to content

Commit a389863

Browse files
MrGevadominicshanshan
authored andcommitted
[None][fix] AD test_trtllm_bench to use small model config and skip loading weights (NVIDIA#8149)
Signed-off-by: Eran Geva <19514940+MrGeva@users.noreply.github.com>
1 parent 12e1f73 commit a389863

File tree

1 file changed

+6
-0
lines changed

1 file changed

+6
-0
lines changed

tests/unittest/_torch/auto_deploy/_utils_test/_model_test_utils.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -452,6 +452,12 @@ def apply_rotary_pos_emb_ds(q, k, cos, sin, position_ids, unsqueeze_dim=1):
452452
"num_hidden_layers": 2,
453453
},
454454
},
455+
"TinyLlama/TinyLlama-1.1B-Chat-v1.0": {
456+
"llm_models_subdir": "llama-models-v2/TinyLlama-1.1B-Chat-v1.0",
457+
"model_kwargs": {
458+
"num_hidden_layers": 2,
459+
},
460+
},
455461
}
456462

457463

0 commit comments

Comments
 (0)