Unverified Commit e6da012a authored by Tim Moon's avatar Tim Moon Committed by GitHub
Browse files

[PyTorch] Disable Flash Attention backend in Userbuffers tests (#2399)



Disable Flash attention in Userbuffers tests
Signed-off-by: default avatarTim Moon <tmoon@nvidia.com>
parent 877b7966
......@@ -120,12 +120,14 @@ def _run_layer_with_overlap(
os.environ["PYTORCH_JIT"] = "0"
os.environ["NVTE_TORCH_COMPILE"] = "0"
os.environ["NVTE_ALLOW_NONDETERMINISTIC_ALGO"] = "0"
os.environ["NVTE_FLASH_ATTN"] = "0"
result = subprocess.run(test_cmd, env=os.environ, capture_output=True, check=False)
os.unsetenv("PYTORCH_JIT")
os.unsetenv("NVTE_TORCH_COMPILE")
os.unsetenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO")
os.unsetenv("NVTE_FLASH_ATTN")
if (
result.returncode != 0
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment