Skip to content

[contrib][DistributedFusedLamb] replace apex.parallel.flat_dist_call with coalescing manager in test #1913

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
crcrpar merged 1 commit into from
Jun 26, 2025
Merged

[contrib][DistributedFusedLamb] replace apex.parallel.flat_dist_call with coalescing manager in test #1913

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 12 additions & 2 deletions apex/contrib/test/optimizers/test_distributed_fused_lamb.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,12 +1,22 @@
import os
import inspect

import torch
from torch.cuda.amp import GradScaler
from torch.testing._internal import common_utils
from apex.parallel.distributed import flat_dist_call
from torch.distributed.distributed_c10d import _coalescing_manager

from apex.contrib.optimizers.distributed_fused_lamb import DistributedFusedLAMB
from apex.transformer.testing.distributed_test_base import NcclDistributedTestBase


def flat_dist_call(param_list: list[torch.Tensor], op, args):
with _coalescing_manager(async_ops=True) as cm:
for p in param_list:
op(p, *args)

cm.wait()


def get_init_weights_func():
@torch.no_grad()
def init_weights(m):
Expand Down