learning/annotated_deep_learning_paper_implementations
1
0
Fork 0
mirror of https://github.com/labmlai/annotated_deep_learning_paper_implementations.git synced 2025-08-16 10:51:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

labml update

Browse Source
This commit is contained in:
Varuna Jayasiri
2023-08-10 19:37:55 +05:30
parent 753fbd3446
commit ec3c7366b7
1 changed files with 3 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
labml_nn/scaling/zero3/finetune_neox.py
View File

@ -80,8 +80,9 @@ def main(rank: int, world_size: int, init_method: str = 'tcp://localhost:23456')
torch.cuda.set_device(device) torch.cuda.set_device(device)
# Create the experiment # Create the experiment
experiment.create(name='zero3_neox', writers={'screen', 'labml'}) experiment.create(name='zero3_neox', writers={'screen', 'labml'},
experiment.distributed(rank, world_size) distributed_world_size=world_size,
distributed_rank=rank)
# Create configurations # Create configurations
conf = Configs() conf = Configs()