cleanup after example

deepspeedai · Feb 8, 2025 · c52cd72 · c52cd72
1 parent 92031f8
commit c52cd72
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/deepspeed/runtime/data_pipeline/data_sampling/variable_batch_size_and_lr_example.py b/deepspeed/runtime/data_pipeline/data_sampling/variable_batch_size_and_lr_example.py
@@ -125,8 +125,8 @@ def to_layers(self):
                 "pin_memory": False,
                 "dynamic_batching": {
                     "enabled": True,
-                    # Path to load the sequence lengths from, as {metrics_path}/seqlen/seqlen_sample_to_metric.bin and *.idx
-                    # If these 2 files dont exist, they'll be output there on the first run, and loaded on subsequent runs.
+                    # Files to load the sequence lengths from: {metrics_path}/seqlen/seqlen_sample_to_metric.bin and *.idx
+                    # If files dont exist, they'll be computed and saved on the first run, and loaded on subsequent runs.
                     "metrics_path": "./curriculum_output/",
                     "lr_scaling_method": "linear",
                     "min_batch_size": 1,
@@ -176,3 +176,5 @@ def to_layers(self):
 
             if engine.data_parallel_group.rank() == 0:
                 print(f"epoch {epoch}, batch {batch_id}, loss {loss.item()}, LRs {lr_scheduler.get_lr()}")
+    dist.barrier()
+    dist.destroy_process_group()