github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-12-02 19:56:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

[fx/meta/rpc] move _meta_registration.py to fx folder / register fx functions with compatibility checks / remove color debug (#1710)

Browse Source 
* [fx] move meta registration

* [fx] fix tests.

* [fx] fix test.

* [fx] fix.

* [meta] refactor meta registration.py.

* [fx] add compatibility descriptions.

* [fx] polish import.

* [fx] add a decorator.

* [fx] fix tests.

* [fx] remove print.

* [fx] edit raise error.

* [fx] edit raise error.

* [fx] add type hint.

* [fx] fix import in experimental.

* [rpc] remove color debug.

* [meta] fix naming.
This commit is contained in:
Super Daniel
2022-10-18 10:44:23 +08:00
committed by GitHub
parent e8d8eda5e7
commit 393f594051
32 changed files with 351 additions and 310 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
tests/test_pipeline/rpc_test_utils.py
View File

@@ -1,19 +1,17 @@
import os
import argparse
import os
import warnings
import torch
from torch import nn
import torch.multiprocessing as mp
import torch.distributed.rpc as rpc
from torch.optim import SGD, Adam, RMSprop, Optimizer
from torch._C._distributed_rpc import _is_current_rpc_agent_set
import torch.distributed as dist
from colorama import Back, Style
from colossalai.pipeline.pipeline_process_group import ppg
from colossalai.logging import disable_existing_loggers
import torch.distributed.rpc as rpc
import torch.multiprocessing as mp
from colossalai import launch
from colossalai.logging import disable_existing_loggers
from colossalai.pipeline.pipeline_process_group import ppg
from torch import nn
from torch._C._distributed_rpc import _is_current_rpc_agent_set
from torch.optim import SGD, Adam, Optimizer, RMSprop
rpc_is_initialized = _is_current_rpc_agent_set