From ec3c7366b7c7a763b0c06a6f703ded6bc2267638 Mon Sep 17 00:00:00 2001 From: Varuna Jayasiri Date: Thu, 10 Aug 2023 19:37:55 +0530 Subject: [PATCH] labml update --- labml_nn/scaling/zero3/finetune_neox.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/labml_nn/scaling/zero3/finetune_neox.py b/labml_nn/scaling/zero3/finetune_neox.py index d83312cd..acd7aeb5 100644 --- a/labml_nn/scaling/zero3/finetune_neox.py +++ b/labml_nn/scaling/zero3/finetune_neox.py @@ -80,8 +80,9 @@ def main(rank: int, world_size: int, init_method: str = 'tcp://localhost:23456') torch.cuda.set_device(device) # Create the experiment - experiment.create(name='zero3_neox', writers={'screen', 'labml'}) - experiment.distributed(rank, world_size) + experiment.create(name='zero3_neox', writers={'screen', 'labml'}, + distributed_world_size=world_size, + distributed_rank=rank) # Create configurations conf = Configs() -- GitLab