| 268 | |
| 269 | # Verify that test is valid |
| 270 | def validate_test(model_w_task, dtype, enable_cuda_graph, enable_triton): |
| 271 | model, task = model_w_task |
| 272 | msg = "" |
| 273 | if enable_cuda_graph and (torch_info["cuda_version"] == "0.0"): |
| 274 | msg = "CUDA not detected, cannot use CUDA Graph" |
| 275 | elif enable_cuda_graph and pkg_version.parse(torch.__version__) < pkg_version.parse("1.10"): |
| 276 | msg = "CUDA Graph is only available in torch versions >= 1.10" |
| 277 | elif "gpt-j-6b" in model: |
| 278 | if dtype != torch.half: |
| 279 | msg = f"Not enough GPU memory to run {model} with dtype {dtype}" |
| 280 | elif enable_cuda_graph: |
| 281 | msg = f"Not enough GPU memory to run {model} with CUDA Graph enabled" |
| 282 | elif "gpt-neox-20b" in model: # TODO: remove this when neox issues resolved |
| 283 | msg = "Skipping gpt-neox-20b for now" |
| 284 | elif ("gpt-neox-20b" in model) and (dtype != torch.half): |
| 285 | msg = f"Not enough GPU memory to run {model} with dtype {dtype}" |
| 286 | elif ("bloom" in model) and (dtype != torch.half): |
| 287 | msg = f"Bloom models only support half precision, cannot use dtype {dtype}" |
| 288 | elif (model not in _bert_models + _roberta_models) and enable_cuda_graph: |
| 289 | msg = "Non bert/roberta models do no support CUDA Graph" |
| 290 | elif enable_triton and not (dtype in [torch.half]): |
| 291 | msg = "Triton is for fp16" |
| 292 | elif enable_triton and not deepspeed.HAS_TRITON: |
| 293 | msg = "triton needs to be installed for the test" |
| 294 | elif (model not in _bert_models + _roberta_models) and enable_triton: |
| 295 | msg = "Triton kernels do not support Non bert/roberta models yet" |
| 296 | |
| 297 | # These should be removed once we fix several inference tests failing |
| 298 | if model in [ |
| 299 | "EleutherAI/pythia-70m-deduped", "distilbert/distilbert-base-cased-distilled-squad", "EleutherAI/gpt-j-6b" |
| 300 | ]: |
| 301 | msg = "Test is currently broken" |
| 302 | return msg |
| 303 | |
| 304 | |
| 305 | @pytest.mark.inference |