huggingface / optimum-neuron

Easy, fast and very cheap training and inference on AWS Trainium and Inferentia chips.

Apache License 2.0

196 stars 59 forks source link

Fix MPMD detected error during training with TP #648

Closed michaelbenayoun closed 2 months ago

michaelbenayoun commented 3 months ago

What does this PR do?

This PR tries to fix the issue reported here https://github.com/aws-neuron/neuronx-distributed/issues/24.

It seems to be linked to torch.autocast:

Before this PR, when it was supposed to be disabled we would do: with torch.autocast(enabled=False):
Now we do: with contexlib.nullcontex():.

For some reason the first approach was leading to different cores trying to execute different graphs, leading to the MPMD detected error.

HuggingFaceDocBuilderDev commented 3 months ago

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.