72 lines
2.3 KiB
Python
72 lines
2.3 KiB
Python
import time
|
|
from copy import deepcopy
|
|
|
|
import colossalai
|
|
import torch
|
|
from colossalai.shardformer import ShardConfig, ShardFormer
|
|
from colossalai.testing import spawn
|
|
|
|
from opensora.acceleration.shardformer.policy.t5_encoder import T5EncoderPolicy
|
|
from opensora.models.text_encoder.t5 import T5Embedder
|
|
|
|
|
|
def run_t5_encoder(rank, world_size, port):
|
|
colossalai.launch({}, rank=rank, world_size=world_size, port=port, host="localhost")
|
|
|
|
# t5 embedder
|
|
t5_path = "./pretrained_models/t5_ckpts"
|
|
hf_t5 = T5Embedder(device="cuda", local_cache=True, cache_dir=t5_path, torch_dtype=torch.float)
|
|
sf_t5 = deepcopy(hf_t5)
|
|
|
|
# create huggingface model as normal
|
|
shard_config = ShardConfig(
|
|
tensor_parallel_process_group=None,
|
|
pipeline_stage_manager=None,
|
|
enable_tensor_parallelism=False,
|
|
enable_fused_normalization=False,
|
|
enable_flash_attention=False,
|
|
enable_jit_fused=True,
|
|
enable_sequence_parallelism=False,
|
|
enable_sequence_overlap=False,
|
|
)
|
|
shard_former = ShardFormer(shard_config=shard_config)
|
|
sharded_model, _ = shard_former.optimize(sf_t5.model, policy=T5EncoderPolicy())
|
|
sf_t5.model = sharded_model
|
|
|
|
# test t5 embedder
|
|
texts = ["Who is the best player in the history of NBA?", "How to study computer science?"]
|
|
for i in range(5):
|
|
hf_embs, hf_masks = hf_t5.get_text_embeddings(texts)
|
|
sf_embs, sf_masks = sf_t5.get_text_embeddings(texts)
|
|
|
|
# check accuracy
|
|
assert torch.allclose(hf_embs, sf_embs, rtol=1e-4, atol=1e-5), f"{hf_embs} \nvs\n{sf_embs}"
|
|
assert torch.allclose(hf_masks, sf_masks), f"{hf_masks} \nvs\n{sf_masks}"
|
|
|
|
# measure perf
|
|
torch.cuda.synchronize()
|
|
hf_start = time.time()
|
|
for i in range(20):
|
|
hf_embs, hf_masks = hf_t5.get_text_embeddings(texts)
|
|
torch.cuda.synchronize()
|
|
hf_end = time.time()
|
|
|
|
# convert sf to fp16
|
|
hf_t5.model = hf_t5.model.half()
|
|
torch.cuda.synchronize()
|
|
sf_start = time.time()
|
|
for i in range(20):
|
|
hf_embs, hf_masks = hf_t5.get_text_embeddings(texts)
|
|
torch.cuda.synchronize()
|
|
sf_end = time.time()
|
|
|
|
print(f"[Performance] native: {hf_end - hf_start}s, shardformer: {sf_end - sf_start} s")
|
|
|
|
|
|
def test_t5_encoder():
|
|
spawn(run_t5_encoder)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_t5_encoder()
|