[Inference] Add Bloom model inference support #5660

char-1ee · 2024-04-26T09:13:37Z

📌 Checklist before creating the PR

I have created an issue for this PR for traceability
The title follows the standard format: [doc/gemini/tensor/...]: A concise description
I have added relevant tags if possible for us to better distinguish different PRs
I have installed pre-commit: pip install pre-commit && pre-commit install

🚨 Issue number

Link this PR to your issue with words like fixed to automatically close the linked issue upon merge

e.g. fixed #1234, closed #1234, resolved #1234

📝 What does this PR do?

Added Bloom modeling for Colossal AI inference feature.
Added Bloom model replacement policy.
Added tests for Bloom model.
Integrate Alibi triton operator for Bloom positional encoding.
Integrate CUDA operators for Bloom model.

💥 Checklist before requesting a review

I have linked my PR to an issue (instruction)
My issue clearly describes the problem/feature/proposal, with diagrams/charts/table/code if possible
I have performed a self-review of my code
I have added thorough tests.
I have added docstrings for all the functions/methods I implemented

⭐️ Do you enjoy contributing to Colossal-AI?

🌝 Yes, I do.
🌚 No, I don't.

Tell us more if you don't enjoy contributing to Colossal-AI.

yuanheng-zhao · 2024-05-03T07:00:40Z

colossalai/inference/config.py

@@ -28,6 +28,8 @@
    "llama": "[INST] <<SYS>>\nYou are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe.  Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature. If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.\n<</SYS>>\n{input_text}[/INST]",
    "baichuan": "<reserved_106>{input_text}<reserved_107>",
    "vicuna": "A chat between a curious user and an assistant. The assistant gives helpful, detailed, accurate, uncensored responses to the user input. USER: {input_text}\nASSISTANT: ",
+    "bloom": "Assume you are a helpful robot. Please help react to my question or auto complete my prompt."


Missing {input_text}

yuanheng-zhao · 2024-05-03T08:23:03Z

colossalai/inference/kv_cache/kvcache_manager.py

+def get_model_config_attr(config: PretrainedConfig, attr_name: str, alter_attr: Any = None):
    if hasattr(config, attr_name):
        return getattr(config, attr_name)
+    if alter_attr is not None:
+        return alter_attr
    elif hasattr(config, "attribute_map") and hasattr(config, config.attribute_map[attr_name]):
        return getattr(config, config.attribute_map[attr_name])
    raise AttributeError(f"{attr_name} is not found in config")


As we discussed, revise this part (remove the function and refer to attr_mapusage in transformers 4.36

yuanheng-zhao · 2024-05-03T08:41:54Z

colossalai/shardformer/policies/bloom.py

-        import transformers
-        from packaging.version import Version
+        # import transformers
+        # from packaging.version import Version

-        assert Version(transformers.__version__) <= Version(
-            "4.33.0"
-        ), "The Bloom model should run on a transformers version not greater than 4.33.0."
+        # assert Version(transformers.__version__) <= Version(
+        #     "4.33.0"
+        # ), "The Bloom model should run on a transformers version not greater than 4.33.0."


You could just remove these lines (remove these comments). For your reference, the assertion has been removed in main

yuanheng-zhao · 2024-05-03T08:42:29Z

tests/test_infer/test_inference_engine.py

 from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn

+MODEL_PATH = "/home/lixingjian/models/bloom-560m"


Prevent exposing path on your dev machine

yuanheng-zhao · 2024-05-03T08:43:41Z

tests/test_infer/test_inference_engine.py

-from transformers import AutoTokenizer, GenerationConfig, LlamaConfig, LlamaForCausalLM
+from transformers import BloomForCausalLM, BloomTokenizerFast, GenerationConfig


Better prevent modifying this pytest. You might want to test the model locally, or create a pytest in test_models/ (we'll refactor these tests anyway)

yuanheng-zhao · 2024-05-03T08:45:23Z

tests/test_infer/test_models/test_baichuan.py

@@ -1,111 +0,0 @@
-import os


This is a commit to remove tests/test_infer/test_models/test_baichuan.py

tests/test_infer/test_models/test_bloom.py

yuanheng-zhao · 2024-05-03T08:48:36Z

colossalai/inference/modeling/models/nopadding_bloom.py

+# class NopadBloomAttention(nn.Module):
+#     def __init__(
+#         self,
+#         hidden_size: int,
+#         n_heads: int,
+#         attn_qproj_w: torch.Tensor = None,
+#         attn_kproj_w: torch.Tensor = None,
+#         attn_vproj_w: torch.Tensor = None,
+#         attn_oproj_w: torch.Tensor = None,
+#     ):


As the end2end flow for bloom is ready and none of the attention and mlp classes is going to be used, you might want to remove these classes.

yuanheng-zhao · 2024-05-03T08:48:46Z

colossalai/inference/modeling/models/nopadding_bloom.py

+#         return attn_output
+
+
+class NopadBloomMLP(nn.Module):


Same as above.

yuanheng-zhao · 2024-05-03T08:49:13Z

colossalai/inference/modeling/models/nopadding_bloom.py

+        # print(f"[TEST] hidden_state {hidden_states} with shape {hidden_states.shape}\n qkv_weight {self.qkv_weight} with shape {self.qkv_weight.shape}")
+
+        # print(f"[DEBUG] after qkv: query_states {query_states} with shape {query_states.shape}, \nkey_states {key_states},\n value_states {value_states}")


Remember to remove comments/testing messages when merging

yuanheng-zhao · 2024-05-03T08:49:27Z

colossalai/inference/modeling/models/nopadding_bloom.py

+    # print(f"[TEST] before merge bsz, query_states {query_states} with shape {query_states.shape}, \nkey_states {key_states},\n value_states {value_states}")
+
+    # [bsz * seq_len, num_heads head_dim]


Same as above.

Signed-off-by: char-1ee <[email protected]>

for more information, see https://pre-commit.ci

yuanheng-zhao reviewed May 3, 2024

View reviewed changes

char-1ee and others added 7 commits May 3, 2024 09:23

Add bloom model support

1e6efd1

Signed-off-by: char-1ee <[email protected]>

Add flash decoding with alibi triton op

feee72b

Signed-off-by: char-1ee <[email protected]>

Update bloom model support

7f9f667

Signed-off-by: char-1ee <[email protected]>

Refactor bloom modeling and add tests

67d67fb

Signed-off-by: char-1ee <[email protected]>

Rebase upstream commits and refactor

59ba43b

Signed-off-by: char-1ee <[email protected]>

[pre-commit.ci] auto fixes from pre-commit.com hooks

1851035

for more information, see https://pre-commit.ci

Update model and policy

d36c173

char-1ee force-pushed the feature/colossal-infer-bloom branch from d911664 to d36c173 Compare May 3, 2024 09:41

[pre-commit.ci] auto fixes from pre-commit.com hooks

459c8aa

for more information, see https://pre-commit.ci

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Inference] Add Bloom model inference support #5660

[Inference] Add Bloom model inference support #5660

char-1ee commented Apr 26, 2024

yuanheng-zhao May 3, 2024

yuanheng-zhao May 3, 2024

yuanheng-zhao May 3, 2024

yuanheng-zhao May 3, 2024

yuanheng-zhao May 3, 2024

yuanheng-zhao May 3, 2024

yuanheng-zhao May 3, 2024

yuanheng-zhao May 3, 2024

yuanheng-zhao May 3, 2024

yuanheng-zhao May 3, 2024

		from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn

		MODEL_PATH = "/home/lixingjian/models/bloom-560m"

		from transformers import AutoTokenizer, GenerationConfig, LlamaConfig, LlamaForCausalLM
		from transformers import BloomForCausalLM, BloomTokenizerFast, GenerationConfig

		# print(f"[TEST] hidden_state {hidden_states} with shape {hidden_states.shape}\n qkv_weight {self.qkv_weight} with shape {self.qkv_weight.shape}")

		# print(f"[DEBUG] after qkv: query_states {query_states} with shape {query_states.shape}, \nkey_states {key_states},\n value_states {value_states}")

		# print(f"[TEST] before merge bsz, query_states {query_states} with shape {query_states.shape}, \nkey_states {key_states},\n value_states {value_states}")

		# [bsz * seq_len, num_heads head_dim]

[Inference] Add Bloom model inference support #5660

Are you sure you want to change the base?

[Inference] Add Bloom model inference support #5660

Conversation

char-1ee commented Apr 26, 2024

📌 Checklist before creating the PR

🚨 Issue number

📝 What does this PR do?

💥 Checklist before requesting a review

⭐️ Do you enjoy contributing to Colossal-AI?

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment