Better error message if workers fall out of sync
Summary:
This kind of issue should be rare, but the exception that was thrown before ("UnpicklingError: invalid load key") was very opaque, so let's use something a bit clearer.
Pull Request resolved:
@@ -148,6 +148,7 @@ def all_gather_list(data, group=None, max_size=16384):
all_reduce(buffer, group=group)
+ try:
result = []
for i in range(world_size):
out_buffer = buffer[i * max_size : (i + 1) * max_size]
@@ -157,3 +158,13 @@ def all_gather_list(data, group=None, max_size=16384):
pickle.loads(bytes(out_buffer[2:size+2].tolist()))
)
return result
+ except pickle.UnpicklingError as e:
+ raise Exception(
+ 'Unable to unpickle data from other workers. all_gather_list requires all '
+ 'workers to enter the function together, so this error usually indicates '
+ 'that the workers have fallen out of sync somehow. Workers can fall out of '
+ 'sync if one of them runs out of memory, or if there are other conditions '
+ 'in your training script that can cause one worker to finish an epoch '
+ 'while other workers are still iterating over their portions of the data.'
+ )
+
