[AsyncTP] Fixes AsyncMM

[fegin]

Stack from ghstack (oldest at bottom):

• [AsyncTP] Use assertEqual instead of allClose for bf16 tests #162041
• -> [AsyncTP] Fixes AsyncMM #162040

The original implementation set beta to be 1, which cause the out (C) being added to the the output. Thus if the output is not initialized as zero beforehand, the output can be incorrect.

Removing the alpha and beta fixes the issue.

Thanks @ngimel to figure out the root cause.

cc @H-Huang @awgu @wanchaol @fduwjj @wz337 @wconstab @d4l3k @pragupta @ezyang @msaroufim

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/162040

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 063094d with merge base 1e0656f ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[ngimel]

If in the tests you set

torch.utils.deterministic.fill_uninitialized_memory=True
torch.use_deterministic_algorithms(True)

you should be able to reliably repro the issue

[fegin]

@ngimel You are right, I can always reproduce the error with the two options you mentioned. But that requires CUBLAS_WORKSPACE_CONFIG=:4096:8 to be exported before running the test. So I guess we cannot add that to the thest.

[ngimel]

huh cc @eqy, exporting this CUBLAS_WORKSPACE_CONFIG=:4096:8 env var is not longer necessary for deterministic mode. Are you getting any errors without it? If it's just warnings you can ignore.

[fegin]

@ngimel I get a hard error. So unfortunately, I cannot ignore it :(

[ngimel]

Ok, this requirement will be removed once #161749 is landed

[ngimel]

You probably can do torch.set_deterministic_debug_mode("warn") + torch.utils.deterministic.fill_uninitialized_memory=True and that way you'd still get the tests to fail, but it won't require env variable even today

[fegin]

You probably can do torch.set_deterministic_debug_mode("warn") + torch.utils.deterministic.fill_uninitialized_memory=True and that way you'd still get the tests to fail, but it won't require env variable even today

This works. Add this to the test and landing the PR.

[fegin]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here