[inference] chatglm2 infer demo (#4724)

* add chatglm2

* add

* gather needed kernels

* fix some bugs

* finish context forward

* finish context stage

* fix

* add

* pause

* add

* fix bugs

* finish chatglm

* fix bug

* change some logic

* fix bugs

* change some logics

* add

* add

* add

* fix

* fix tests

* fix

tests/kit/model_zoo/transformers/chatglm2.py

@@ -39,6 +39,21 @@ config = ChatGLMConfig(
     padded_vocab_size=65024,
     hidden_size=64,
     num_attention_heads=8,
+    kv_channels=16,
+    rmsnorm=True,
+    original_rope=True,
+    use_cache=True,
+    torch_dtype=torch.float32,
+)
+
+infer_config = ChatGLMConfig(
+    num_layers=2,
+    padded_vocab_size=65024,
+    hidden_size=128,
+    num_attention_heads=8,
+    multi_query_attention=True,
+    multi_query_group_num=2,
+    kv_channels=16,
     rmsnorm=True,
     original_rope=True,
     use_cache=True,

tests/test_infer/test_chatglm2_infer.py

@@ -0,0 +1,73 @@
+import os
+
+import pytest
+import torch
+import torch.distributed as dist
+from packaging import version
+from transformers import AutoTokenizer
+
+import colossalai
+from colossalai.inference.tensor_parallel.engine import TPInferEngine
+from colossalai.logging import disable_existing_loggers
+from colossalai.shardformer import ShardConfig
+from colossalai.shardformer.modeling.chatglm2_6b.modeling_chatglm import ChatGLMForConditionalGeneration
+from colossalai.testing import clear_cache_before_run, parameterize, rerun_if_address_is_in_use, spawn
+from tests.kit.model_zoo.transformers.chatglm2 import infer_config
+
+os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "true"
+TPSIZE = 1
+BATCH_SIZE = 8
+MAX_INPUT_LEN = 12
+MAX_OUTPUT_LEN = 100
+CUDA_SUPPORT = version.parse(torch.version.cuda) > version.parse("11.5")
+
+
+@parameterize(
+    "test_config",
+    [
+        {
+            "tp_size": TPSIZE,
+        }
+    ],
+)
+def run_chatglm2_test(test_config):
+    tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True)
+    # pad_token_id = 0
+    model_fn = lambda: ChatGLMForConditionalGeneration(infer_config, empty_init=False)
+    orig_model = model_fn()
+    orig_model = orig_model.half()
+    text = ["how is the weather today?"]
+    input_ids = tokenizer.batch_encode_plus(text, return_tensors="pt", padding=True)
+    shard_config = ShardConfig(
+        enable_tensor_parallelism=True if test_config["tp_size"] > 1 else False, inference_only=True
+    )
+    infer_engine = TPInferEngine(orig_model, shard_config, BATCH_SIZE, MAX_INPUT_LEN, MAX_OUTPUT_LEN)
+
+    generate_kwargs = dict(max_new_tokens=MAX_OUTPUT_LEN, do_sample=False)
+    outputs = infer_engine.generate(input_ids, **generate_kwargs)
+    assert outputs is not None
+
+    # print("outputs.shape: ", outputs[0].shape)
+    # print("outputs: ", outputs[0])
+    if not dist.is_initialized() or dist.get_rank() == 0:
+        for o in outputs:
+            output_text = tokenizer.decode(o)
+            print(output_text)
+
+
+def check_chatglm2(rank, world_size, port):
+    disable_existing_loggers()
+    colossalai.launch(config={}, rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl")
+    run_chatglm2_test()
+
+
+@pytest.mark.skipif(not CUDA_SUPPORT, reason="kv-cache manager engine requires cuda version to be higher than 11.5")
+@pytest.mark.dist
+@rerun_if_address_is_in_use()
+@clear_cache_before_run()
+def test_chatglm2():
+    spawn(check_chatglm2, TPSIZE)
+
+
+if __name__ == "__main__":
+    test_chatglm2()

tests/test_infer_ops/triton/test_llama2_token_attn.py

@@ -0,0 +1,65 @@
+import pytest
+import torch
+from packaging import version
+
+try:
+    pass
+
+    from colossalai.kernel.triton.token_attention_kernel import Llama2TokenAttentionForwards
+
+    HAS_TRITON = True
+except ImportError:
+    HAS_TRITON = False
+    print("please install triton from https://github.com/openai/triton")
+
+TRITON_CUDA_SUPPORT = version.parse(torch.version.cuda) > version.parse("11.4")
+
+
+def torch_att(xq, xk, xv, bs, seqlen, num_head, head_dim):
+    xq = xq.view(bs, 1, num_head, head_dim)
+    xk = xk.view(bs, seqlen, num_head, head_dim)
+    xv = xv.view(bs, seqlen, num_head, head_dim)
+
+    logics = torch.sum(xq * xk, dim=3, keepdim=False) * 1 / (head_dim**0.5)
+    prob = torch.softmax(logics, dim=1)
+    prob = prob.view(bs, seqlen, num_head, 1)
+
+    return torch.sum(prob * xv, dim=1, keepdim=False)
+
+
+@pytest.mark.skipif(
+    not TRITON_CUDA_SUPPORT or not HAS_TRITON, reason="triton requires cuda version to be higher than 11.4"
+)
+def test():
+    Z, head_num, seq_len, head_dim = 2, 32, 2048, 128
+    dtype = torch.float16
+
+    # attn out: 2,4096
+    q = torch.empty((Z, head_num, head_dim), dtype=dtype, device="cuda").normal_(mean=0.1, std=0.2)
+    k = torch.empty((Z * seq_len, head_num, head_dim), dtype=dtype, device="cuda").normal_(mean=0.4, std=0.2)
+    v = torch.empty((Z * seq_len, head_num, head_dim), dtype=dtype, device="cuda").normal_(mean=0.3, std=0.2)
+    o = torch.empty_like()
+    # o = torch.empty((Z, head_num, head_dim), dtype=dtype, device="cuda").normal_(mean=0.3, std=0.2)
+
+    max_kv_cache_len = seq_len
+    kv_cache_start_loc = torch.zeros((Z,), dtype=torch.int32, device="cuda")
+    kv_cache_loc = torch.zeros((Z, seq_len), dtype=torch.int32, device="cuda")
+    kv_cache_seq_len = torch.ones((Z,), dtype=torch.int32, device="cuda")
+    other_kv_index = 2048
+
+    kv_cache_seq_len[:] = seq_len
+    kv_cache_start_loc[0] = 0
+    kv_cache_start_loc[1] = seq_len
+
+    for i in range(Z):
+        kv_cache_loc[i, :] = torch.arange(i * seq_len, (i + 1) * seq_len, dtype=torch.int32, device="cuda")
+
+    Llama2TokenAttentionForwards.token_attn(
+        q, k, v, o, kv_cache_loc, kv_cache_start_loc, kv_cache_seq_len, max_kv_cache_len, other_kv_index
+    )
+    torch_out = torch_att(q, k, v, Z, seq_len, head_num, head_dim)
+    assert torch.allclose(torch_out, o, atol=1e-3, rtol=0)
+
+
+if __name__ == "__main__":
+    test()