github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-07-16 00:17:55 +00:00
Code Issues Packages Projects Releases Wiki Activity

[test] fixed the activation codegen test (#1447)

Browse Source 
* [test] fixed the activation codegen test

* polish code
This commit is contained in:
Frank Lee 2022-08-12 14:52:31 +08:00 committed by GitHub
parent 0f3042363c
commit 5a52e21fe3
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 15 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
colossalai/utils/common.py
View File

@ -1,7 +1,6 @@
#!/usr/bin/env python
# -*- encoding: utf-8 -*-
import os
from pprint import pp
import random
import socket
from pathlib import Path

21
tests/test_fx/test_codegen/test_activation_checkpoint_codegen.py
View File

@ -1,6 +1,7 @@
from operator import mod
import torch
import pytest
import torch.multiprocessing as mp
from torch.utils.checkpoint import checkpoint
from torch.fx import GraphModule
from colossalai.fx import ColoTracer
@ -42,10 +43,9 @@ class MyModule(torch.nn.Module):
return y1 + y2 + y3 + y4
@pytest.mark.skipif(not with_codegen, reason='torch version is lower than 1.12.0')
def test_act_ckpt_codegen():
def _run_act_ckpt_codegen(rank):
# launch colossalai to make sure we could execute colossalai.utils.checkpoint currectly
colossalai.launch(config={}, rank=0, world_size=1, host='localhost', port=free_port(), backend='nccl')
colossalai.launch(config={}, rank=rank, world_size=1, host='localhost', port=free_port(), backend='nccl')
# build model and run forward
model = MyModule()
@ -90,10 +90,14 @@ def test_act_ckpt_codegen():
gpc.destroy()
@pytest.mark.skipif(with_codegen, reason='torch version is equal to or higher than 1.12.0')
def test_act_ckpt_python_code_torch11():
@pytest.mark.skipif(not with_codegen, reason='torch version is lower than 1.12.0')
def test_act_ckpt_codegen():
mp.spawn(_run_act_ckpt_codegen, nprocs=1)
def _run_act_ckpt_python_code_torch11(rank):
# launch colossalai to make sure we could execute colossalai.utils.checkpoint currectly
colossalai.launch(config={}, rank=0, world_size=1, host='localhost', port=free_port(), backend='nccl')
colossalai.launch(config={}, rank=rank, world_size=1, host='localhost', port=free_port(), backend='nccl')
# build model and run forward
model = MyModule()
@ -138,6 +142,11 @@ def test_act_ckpt_python_code_torch11():
gpc.destroy()
@pytest.mark.skipif(with_codegen, reason='torch version is equal to or higher than 1.12.0')
def test_act_ckpt_python_code_torch11():
mp.spawn(_run_act_ckpt_python_code_torch11, nprocs=1)
if __name__ == '__main__':
test_act_ckpt_codegen()