[Pipeline] Drop last batch in DeepSpeed scripts (awslabs#19)

szhengac · Jan 27, 2023 · e106607 · e106607
1 parent 5aaed77
commit e106607
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/examples/data_util.py b/examples/data_util.py
@@ -68,13 +68,15 @@ def get_dataloader(model_name, micro_batch_size, enable_pipeline, cache_dir=None
         batch_size=micro_batch_size,
         sampler=DistributedSampler(train_dataset),
         collate_fn=get_data_move_and_group_fn(enable_pipeline),
+        drop_last=True,
     )
 
     val_loader = DataLoader(
         val_dataset,
         batch_size=micro_batch_size,
         sampler=DistributedSampler(val_dataset),
         collate_fn=get_data_move_and_group_fn(enable_pipeline),
+        drop_last=True,
     )
     return train_loader, val_loader
 

diff --git a/slapo/schedule.py b/slapo/schedule.py
@@ -1106,7 +1106,6 @@ def build(
         sch.metadata.tie_weights = analyze_tie_weights(
             sch.mod, is_pipeline_partitioned=True
         )
-        print(f"tie_weight_groups: {sch.metadata.tie_weights}")
 
     # delay initialization
     if init_weights: