Minor fixes and improvements

Author: Shiqiang Wang

fms_fsdp/utils/dataloader_utils.py

@@ -88,8 +88,9 @@ def get_data_loader(cfg, rank, world_size, postprocess=[causal_lm]):
     ), f"File type {cfg.file_type} is not recognized ({list(_handler_map.keys())})"
     if cfg.file_type == "hf_parquet" or cfg.file_type == "auto":
         filehandler = _handler_map[cfg.file_type](cfg.tokenizer_path, cols)
-    else:
-        filehandler = _handler_map[cfg.file_type, cols]
+    elif cfg.file_type == "arrow":
+        filehandler = _handler_map[cfg.file_type](cols)
+        
     # Base reader layer
     data = StreamingDocDataset(
         cfg.data_path,

fms_fsdp/utils/dataset_utils.py

@@ -939,7 +939,7 @@ def setup(self):
             # listdir, assemble shardfraglist (ind -> shard, frag)
             shards = [
                 os.path.join(root, name)[len(datapath) + 1 :]
-                for root, dirs, files in os.walk(datapath, topdown=False)
+                for root, dirs, files in os.walk(datapath, topdown=False, followlinks=True)
                 for name in files
                 if self.filehandler.is_legal(os.path.join(root, name))
             ]

fms_fsdp/utils/train_utils.py

@@ -43,20 +43,27 @@ def train(
                 import wandb  # type: ignore
             except ImportError:
                 raise ImportError("tracker is set to wandb but wandb is not installed.")
+            # dist.barrier()
             if rank == 0:
-                print(f"--> wandb is enabled!")
+                print(f"--> Started initializing wandb", flush=True)
                 try:
                     wandb.init(
                         project=project_name,
                         dir=tracker_dir,
                         resume="allow",
                         id=run_id,
+                        # mode='offline',
+                        settings=wandb.Settings(
+                                init_timeout=3600,
+                            )
                     )
                 except wandb.errors.UsageError:
                     raise ValueError(
                         "wandb failed to init, did you pass your wandb api key via WANDB_API_KEY?"
                     )
                 wandb.config = asdict(cfg)
+                print(f"--> wandb is enabled!", flush=True)
+            # dist.barrier()
 
         if cfg.tracker == "aim":
             try:
@@ -102,7 +109,7 @@ def train(
         if profiler:
             profiler.step()
 
-        if batch_idx % cfg.report_interval == 0:
+        if batch_idx % cfg.report_interval == 0 or batch_idx == start_step + 1:
             dist.all_reduce(ddp_stats, op=dist.ReduceOp.SUM)
             train_loss = ddp_stats[0] / ddp_stats[2]
             g_norm = ddp_stats[1] / ddp_stats[2]
@@ -146,7 +153,7 @@ def train(
                     "overall token per day:",
                     int(new_tokens_seen / elapsed_time * 3600 * 24),
                 )
-                if cfg.tracker:
+                if cfg.tracker and batch_idx > start_step + 1:
                     vals_to_track = {
                         "learning rate": current_lr,
                         "loss": current_loss,