Xla: Client mesh is unhealthy with uneven heartbeats

[myleott]

I'm doing a large-ish training run (v3-256) with a dataset that takes ~10 minutes to load at the beginning of training. After a few minutes of data loading I'm hitting an error: "Client mesh is unhealthy with uneven heartbeats," after which everything restarts.

It looks like this is happening because mark_step isn't called for more than 300 seconds (based on this timeout).

Any tips on how to handle this? Can I disable the heartbeats while I'm loading data? Alternatively should I be calling mark_step manually?

[taylanbil]

This is being looked at by @jysohn23. For the time being, as a workaround, you can increase the timeout and run your distributed job as

python /path/to/edited/xla_dist.py --... And your new timeout should be in effect.

Curious, how many restarts are you observing? Does it restart a second time after the first restart?

[myleott]

There were 3 restarts before I noticed it and killed the job.

In this particular case the dataset is partitioned into a few pieces, so I was able to get things working by adding mark_step between loading of each partition. But a more general solution would be preferred :)

[jysohn23]

Hi, thanks for reporting the issue.

So I'm not sure how the data loading bits are happening, but it sounds like only the master may be doing the dataloading work? If all the workers were doing the loading of the data, all the workers should be waiting to finish data loading and only then all move forward to mark_step and fall into the 2nd case listed below (but based on the logs it looks like its the former).

We have 2 timeouts for declaring a training run in "hanging"

1. mark_step counts are uneven (one worker may have died) - this one has a shorter timeout (5 min).
2. mark_step counts are all the same for all workers - maybe it's checkpointing/loading data and thus the longer timeout (30 min).

Do you know whether there's a worker that's getting to the first worker first? You should be able to add a barrier with the following feature that @dlibenzi added: https://github.com/pytorch/xla/pull/1655

[myleott]

Ah makes sense. All workers are loading the data, but maybe one of them is finishing much faster than the others, causing an uneven mark_step count.

xmp.rendezvous() looks like a great solution, thanks!

[jysohn23]

Actually location of xmp.rendezvous() has been moved to xm.rendezvous() 😄: https://github.com/pytorch/xla/issues/1654#issuecomment-589390209

[myleott]

This worked, thanks again!

[myleott]

Now I'm running into a race condition of sorts. This is on a v3-32. I had something like this:

load_data()  # ~10 minutes
xm.rendezvous()
xm.mark_step()
begin_training()

The rendezvous call blocks until all workers reach it, but there's still a delay between calling mark_step and having it picked up by all the clients. So if the _check_client_mesh_health function runs before all of the client steps are marked, then we get a "Client mesh is unhealthy with uneven heartbeats" error.

Any suggestions? It seems once uneven_health_timeout seconds have passed since the last mark_step, there's no way to definitively avoid the above race condition. Can rendezvous reset the heartbeat?

[dlibenzi]

Does load_data() do any XLA tensor work?
If not, I don't think the mark_step() is necessary.
Otherwise how about?

load_data()  # ~10 minutes
xm.mark_step()
xm.rendezvous()
begin_training()

[myleott]

Does load_data() do any XLA tensor work?

No, but there may have been XLA work done prior to load_data(), e.g., we often split large datasets into smaller pieces and load them at the start of each "epoch."

Otherwise how about? (...)

I don't think that will work. The first client to call mark_step will immediately trigger the unhealthy mesh warning, right?

[dlibenzi]

In theory the first mark_step() should either find no XLA work, or find a relatively simple graph, so compilation and execution should be fast.
But, @jysohn23 , to unblock this until we investigate further, please make those timeout environment configurable.

[myleott]

I mean, what if you have something like this:

xm.mark_step()
load_data()   # ~10 minutes
xm.mark_step()
xm.rendezvous()

This will cause the unhealthy mesh warning because as soon as any worker calls the second mark_step, we'll be in an unhealthy state.

This happens in practice for me because my training loop looks something like this:

while training:
  load_data()
  xm.rendezvous()
  train_for_one_epoch()

[dlibenzi]

I did not realize that was in a loop. At that point it unrolls to the state you are describing.
We will unblock you with configurable timeouts, and then we will see whether it makes sense to split the minute in case of HW failures detections.