From b7c0d7ceea4825c07823ea862b9f668f814bded0 Mon Sep 17 00:00:00 2001 From: Eitan Turok Date: Fri, 9 Aug 2024 18:14:50 +0000 Subject: [PATCH] fix linting --- megablocks/layers/mpu.py | 1 + 1 file changed, 1 insertion(+) diff --git a/megablocks/layers/mpu.py b/megablocks/layers/mpu.py index 63de336..239f75f 100644 --- a/megablocks/layers/mpu.py +++ b/megablocks/layers/mpu.py @@ -41,6 +41,7 @@ def copy_expert_model_parallel_attributes( getattr(source_tensor, 'expert_model_parallel'), ) + def synchronized_print(group, *x): world_size = torch.distributed.get_world_size(group) rank = torch.distributed.get_rank(group)