mirror of https://github.com/hpcaitech/ColossalAI
62 lines
1.7 KiB
Python
62 lines
1.7 KiB
Python
#!/usr/bin/env python
|
|
# -*- encoding: utf-8 -*-
|
|
|
|
import pytest
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from torch.utils.checkpoint import checkpoint
|
|
|
|
from colossalai.context.parallel_mode import ParallelMode
|
|
from colossalai.context.random import add_seed, seed, set_mode
|
|
from colossalai.utils import checkpoint
|
|
|
|
|
|
def forward(x, weight):
|
|
out = torch.matmul(x, weight)
|
|
with seed(ParallelMode.DATA):
|
|
out_ = F.dropout(out, p=0.4, training=True)
|
|
return out_
|
|
|
|
|
|
@pytest.mark.gpu
|
|
def test_activation_checkpointing():
|
|
add_seed(ParallelMode.GLOBAL, 1024)
|
|
set_mode(ParallelMode.GLOBAL)
|
|
global_cuda_rng_state = torch.cuda.get_rng_state()
|
|
add_seed(ParallelMode.DATA, 1026)
|
|
set_mode(ParallelMode.DATA)
|
|
data_parallel_cuda_rng_state = torch.cuda.get_rng_state()
|
|
set_mode(ParallelMode.GLOBAL)
|
|
|
|
# normal
|
|
data = torch.rand(2, 2, requires_grad=True).cuda()
|
|
data.retain_grad()
|
|
weight = torch.rand(2, 4, requires_grad=True).cuda()
|
|
|
|
data_ = data.clone().detach()
|
|
data_.requires_grad = True
|
|
data_.retain_grad()
|
|
weight_ = weight.clone().detach()
|
|
weight_.requires_grad = True
|
|
|
|
out = forward(data, weight)
|
|
loss = out.sum()
|
|
loss.backward()
|
|
|
|
# checkpoint
|
|
set_mode(ParallelMode.GLOBAL)
|
|
torch.cuda.set_rng_state(global_cuda_rng_state)
|
|
set_mode(ParallelMode.DATA)
|
|
torch.cuda.set_rng_state(data_parallel_cuda_rng_state)
|
|
set_mode(ParallelMode.GLOBAL)
|
|
out = checkpoint(forward, data_, weight_)
|
|
loss = out.sum()
|
|
loss.backward()
|
|
|
|
assert torch.all(data.grad == data_.grad), 'Gradient of the input does not match'
|
|
torch.cuda.empty_cache()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
test_activation_checkpointing()
|