github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity
ColossalAI/tests/test_shardformer/test_model/_utils.py

77 lines
2.9 KiB
Python
Raw Normal View History

[shardformer] adapted T5 and LLaMa test to use kit (#4049) * [shardformer] adapted T5 and LLaMa test to use kit * polish code
2023-06-21 01:32:46 +00:00
import copy
[shardformer] support lazy init (#4202) * [shardformer] support lazy init * [shardformer] linear support lazy init * [shardformer] embedding support lazy init * [shardformer] norm support lazy init * [shardformer] fused linear support lazy init * [test] update shardformer test layer * [test] shardformer with lazy init fit ddp * [lazy] hotfix deepcopy of param * [shardformer] fix bert policy and update test * [shardformer] fix bloom policy and update test * [shardformer] fix opt policy and update test * [shardformer] fix t5 policy and update test * [shardformer] fix gpt2 policy and update test * [shardformer] fix llama policy and update test
2023-07-10 02:48:53 +00:00
from contextlib import nullcontext
[shardformer] adapted T5 and LLaMa test to use kit (#4049) * [shardformer] adapted T5 and LLaMa test to use kit * polish code
2023-06-21 01:32:46 +00:00
[shardformer] support inplace sharding (#4251) * [shardformer] embedding support inplace sharding * [shardformer] linear support inplace sharding * [shardformer] layernorm support inplace sharding * [shardformer] qkv support inplace sharding * [test] update shardformer layer test * [shardformer] fix shared param sharding * [shardformer] fix bert policy * [shardformer] fix bloom policy * [shardformer] fix llama policy * [shardformer] fix opt policy * [shardformer] fix t5 policy * [shardformer] fix fused qkv linear * [shardformer] fix bugs * force sync * [test] fix bugs * [test] fix transformer version
2023-07-20 02:39:06 +00:00
import torch
from torch.nn import Module
[shardformer] support lazy init (#4202) * [shardformer] support lazy init * [shardformer] linear support lazy init * [shardformer] embedding support lazy init * [shardformer] norm support lazy init * [shardformer] fused linear support lazy init * [test] update shardformer test layer * [test] shardformer with lazy init fit ddp * [lazy] hotfix deepcopy of param * [shardformer] fix bert policy and update test * [shardformer] fix bloom policy and update test * [shardformer] fix opt policy and update test * [shardformer] fix t5 policy and update test * [shardformer] fix gpt2 policy and update test * [shardformer] fix llama policy and update test
2023-07-10 02:48:53 +00:00
from colossalai.lazy import LazyInitContext
[shardformer] adapted T5 and LLaMa test to use kit (#4049) * [shardformer] adapted T5 and LLaMa test to use kit * polish code
2023-06-21 01:32:46 +00:00
from colossalai.shardformer import ShardConfig, ShardFormer
[shardformer] support lazy init (#4202) * [shardformer] support lazy init * [shardformer] linear support lazy init * [shardformer] embedding support lazy init * [shardformer] norm support lazy init * [shardformer] fused linear support lazy init * [test] update shardformer test layer * [test] shardformer with lazy init fit ddp * [lazy] hotfix deepcopy of param * [shardformer] fix bert policy and update test * [shardformer] fix bloom policy and update test * [shardformer] fix opt policy and update test * [shardformer] fix t5 policy and update test * [shardformer] fix gpt2 policy and update test * [shardformer] fix llama policy and update test
2023-07-10 02:48:53 +00:00
def build_model(model_fn, enable_fused_normalization=True, enable_tensor_parallelism=True, use_lazy_init: bool = False):
ctx = LazyInitContext() if use_lazy_init else nullcontext()
with ctx:
# create new model
org_model = model_fn()
model_copy = copy.deepcopy(org_model)
if use_lazy_init:
ctx.materialize(org_model)
[shardformer] adapted T5 and LLaMa test to use kit (#4049) * [shardformer] adapted T5 and LLaMa test to use kit * polish code
2023-06-21 01:32:46 +00:00
# shard model
[shardformer] made tensor parallelism configurable (#4144) * [shardformer] made tensor parallelism configurable * polish code
2023-07-04 01:57:03 +00:00
shard_config = ShardConfig(enable_fused_normalization=enable_fused_normalization,
enable_tensor_parallelism=enable_tensor_parallelism)
[shardformer] adapted T5 and LLaMa test to use kit (#4049) * [shardformer] adapted T5 and LLaMa test to use kit * polish code
2023-06-21 01:32:46 +00:00
shard_former = ShardFormer(shard_config=shard_config)
[test] update shardformer tests
2023-07-05 06:30:17 +00:00
sharded_model, shared_params = shard_former.optimize(model_copy)
[shardformer] support lazy init (#4202) * [shardformer] support lazy init * [shardformer] linear support lazy init * [shardformer] embedding support lazy init * [shardformer] norm support lazy init * [shardformer] fused linear support lazy init * [test] update shardformer test layer * [test] shardformer with lazy init fit ddp * [lazy] hotfix deepcopy of param * [shardformer] fix bert policy and update test * [shardformer] fix bloom policy and update test * [shardformer] fix opt policy and update test * [shardformer] fix t5 policy and update test * [shardformer] fix gpt2 policy and update test * [shardformer] fix llama policy and update test
2023-07-10 02:48:53 +00:00
return org_model.cuda(), sharded_model.cuda()
[shardformer] adapted T5 and LLaMa test to use kit (#4049) * [shardformer] adapted T5 and LLaMa test to use kit * polish code
2023-06-21 01:32:46 +00:00
[pipeline] Bert pipeline for shardformer and its tests (#4197) * add pipeline forward * complete pipeline forward check * fix bert forward without pipeline * fix comments * discard useless line * add todo * clean prints * fix distribute layers
2023-07-10 05:58:58 +00:00
def build_pipeline_model(model_fn,
stage_manager=None,
enable_fused_normalization=False,
enable_tensor_parallelism=False,
use_lazy_init: bool = False):
ctx = LazyInitContext() if use_lazy_init else nullcontext()
with ctx:
# create new model
org_model = model_fn()
model_copy = copy.deepcopy(org_model)
if use_lazy_init:
ctx.materialize(org_model)
# shard model
shard_config = ShardConfig(enable_fused_normalization=enable_fused_normalization,
enable_tensor_parallelism=enable_tensor_parallelism,
pipeline_stage_manager=stage_manager)
[pipeline] Llama pipeline (#4205) * bloom policy * llama pipeline forward and tests * fix the output and attention_mask * fix name * bind argument to policy * Revert "bloom policy" This reverts commit 8dee68a0a22568dbeed6d4563372b25e1e825fb0. This policy should be revert and copied to feature/bloom * revert the bloom changes * cancel unneeded inputs * gpt
2023-07-11 03:37:26 +00:00
[pipeline] Bert pipeline for shardformer and its tests (#4197) * add pipeline forward * complete pipeline forward check * fix bert forward without pipeline * fix comments * discard useless line * add todo * clean prints * fix distribute layers
2023-07-10 05:58:58 +00:00
shard_former = ShardFormer(shard_config=shard_config)
sharded_model, shared_params = shard_former.optimize(model_copy)
return org_model.cuda(), sharded_model.cuda()
[shardformer] adapted T5 and LLaMa test to use kit (#4049) * [shardformer] adapted T5 and LLaMa test to use kit * polish code
2023-06-21 01:32:46 +00:00
def run_forward(original_model, sharded_model, data_gen_fn, output_transform_fn, loss_fn):
# prepare input
data = data_gen_fn()
data = {k: v.cuda() for k, v in data.items()}
# switch to train mode
original_model.train()
sharded_model.train()
# run forward
org_output = original_model(**data)
org_output = output_transform_fn(org_output)
org_loss = loss_fn(org_output)
shard_output = sharded_model(**data)
shard_output = output_transform_fn(shard_output)
shard_loss = loss_fn(shard_output)
[shardformer] supported fused normalization (#4112)
2023-06-30 01:32:37 +00:00
return org_output, org_loss, shard_output, shard_loss
[shardformer] support inplace sharding (#4251) * [shardformer] embedding support inplace sharding * [shardformer] linear support inplace sharding * [shardformer] layernorm support inplace sharding * [shardformer] qkv support inplace sharding * [test] update shardformer layer test * [shardformer] fix shared param sharding * [shardformer] fix bert policy * [shardformer] fix bloom policy * [shardformer] fix llama policy * [shardformer] fix opt policy * [shardformer] fix t5 policy * [shardformer] fix fused qkv linear * [shardformer] fix bugs * force sync * [test] fix bugs * [test] fix transformer version
2023-07-20 02:39:06 +00:00
def check_state_dict(org_model: Module, sharded_model: Module, name: str = ''):
org_sd = org_model.state_dict()
shard_sd = sharded_model.state_dict()
for k, v in org_sd.items():
assert k in shard_sd, f'{name} {k} not in sharded model'
shard_v = shard_sd[k]
assert v.shape == shard_v.shape, f'{name} {k} shape mismatch, {v.shape} vs {shard_v.shape}'
assert v.dtype == shard_v.dtype, f'{name} {k} dtype mismatch, {v.dtype} vs {shard_v.dtype}'
assert torch.equal(v, shard_v), f'{name} {k} value mismatch'