address pr comments

llmfoundry/data/finetuning/dataloader.py

@@ -617,13 +617,10 @@ def _download_remote_hf_dataset(remote_path: str, split: str) -> str:
             with open(signal_file_path, 'wb') as f:
                 f.write(b'local_rank0_completed_download')
 
-        print(dist.get_local_rank(), f'signal_file_path: {signal_file_path}')
-
         # Avoid the collective call until the local rank zero has finished trying to download the dataset
         # so that we don't timeout for large downloads. This syncs all processes on the node
         with dist.local_rank_zero_download_and_wait(signal_file_path):
             # Then, wait to ensure every node has finished trying to download the dataset
-            print('GOT TO BARRIER')
             dist.barrier()
 
         # clean up signal file

llmfoundry/utils/file_utils.py

@@ -13,9 +13,12 @@ def dist_mkdtemp() -> str:
         str: The path to the temporary directory.
     """
     tempdir = None
-    if dist.get_local_rank() == 0:
+    local_rank = dist.get_local_rank()
+    global_rank = dist.get_global_rank()
+    if local_rank == 0:
         tempdir = tempfile.mkdtemp()
-    tempdir = dist.all_gather_object(tempdir)[0]
+
+    tempdir = dist.all_gather_object(tempdir)[global_rank - local_rank]
     if tempdir is None:
         raise RuntimeError('Dist operation to get tempdir failed.')
     return tempdir