github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-26 04:03:58 +00:00
Code Issues Packages Projects Releases Wiki Activity

html refactor (#555)

Browse Source
This commit is contained in:
Liang Bowen
2022-03-31 11:36:56 +08:00
committed by GitHub
parent d1211148a7
commit 2c45efc398
133 changed files with 274 additions and 672 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
colossalai/utils/profiler/prof_utils.py
View File

@@ -73,25 +73,26 @@ class ProfilerContext(object):
"""
Profiler context manager
Usage:
::
```python
world_size = 4
inputs = torch.randn(10, 10, dtype=torch.float32, device=get_current_device())
outputs = torch.empty(world_size, 10, 10, dtype=torch.float32, device=get_current_device())
outputs_list = list(torch.chunk(outputs, chunks=world_size, dim=0))
```python
world_size = 4
inputs = torch.randn(10, 10, dtype=torch.float32, device=get_current_device())
outputs = torch.empty(world_size, 10, 10, dtype=torch.float32, device=get_current_device())
outputs_list = list(torch.chunk(outputs, chunks=world_size, dim=0))
cc_prof = CommProfiler()
cc_prof = CommProfiler()
with ProfilerContext([cc_prof]) as prof:
op = dist.all_reduce(inputs, async_op=True)
dist.all_gather(outputs_list, inputs)
op.wait()
dist.reduce_scatter(inputs, outputs_list)
dist.broadcast(inputs, 0)
dist.reduce(inputs, 0)
with ProfilerContext([cc_prof]) as prof:
op = dist.all_reduce(inputs, async_op=True)
dist.all_gather(outputs_list, inputs)
op.wait()
dist.reduce_scatter(inputs, outputs_list)
dist.broadcast(inputs, 0)
dist.reduce(inputs, 0)
prof.show()
```
prof.show()
```
"""
def __init__(self, profilers: List[BaseProfiler] = None, enable: bool = True):