Merge branch 'main' of https://github.com/hpcaitech/ColossalAI

applications/Chat/examples/example_data_reformat.py

-jsonl_file = 'seed_prompts_xx.jsonl'  # seed_prompts_en.jsonl or seed_prompts_ch.json from InstructionWild
-reformat_file = 'prompts_xx.jsonl'  # reformat jsonl file used as Prompt dataset in Stage3
-
-data = ''
-with open(jsonl_file, 'r', encoding="utf-8") as f1:
-    for jsonstr in f1.readlines():
-        jsonstr = '\t' + jsonstr.strip('\n') + ',\n'
-        data = data + jsonstr
-    data = '[\n' + data + ']'
-
-with open(reformat_file, 'w') as f2:
-    f2.write(data)
\ No newline at end of file

applications/Chat/examples/generate_conversation_dataset.py

+import argparse
+import json
+
+from datasets import load_dataset
+
+
+def generate_alpaca():
+    # We can convert dataset with the same format("instruction", "input", "output") as Alpaca into a one-round conversation.
+    conversation_dataset = []
+    dataset = load_dataset("tatsu-lab/alpaca", split="train")
+
+    instructions = dataset["instruction"]
+    inputs = dataset["input"]
+    outputs = dataset["output"]
+
+    assert len(instructions) == len(inputs) == len(outputs)
+
+    for idx in range(len(instructions)):
+        human_utterance = instructions[idx] + "\n\n" + inputs[idx] if inputs[idx] else instructions[idx]
+        human = {"from": "human", "value": human_utterance}
+
+        gpt_utterance = outputs[idx]
+        gpt = {"from": "gpt", "value": gpt_utterance}
+
+        conversation = dict(type="instruction", language="English", dataset="Alpaca", conversations=[human, gpt])
+        conversation_dataset.append(conversation)
+
+    return conversation_dataset
+
+
+def generate_sharegpt():
+    # ShareGPT data requires less processing.
+    conversation_dataset = []
+    dataset = load_dataset(
+        "anon8231489123/ShareGPT_Vicuna_unfiltered",
+        data_files="ShareGPT_V3_unfiltered_cleaned_split_no_imsorry.json",
+        split="train",
+    )
+
+    conversations = dataset["conversations"]
+
+    for idx in range(len(conversations)):
+        for conv in conversations[idx]:
+            # We don't need markdown and text value.
+            del conv["markdown"]
+            del conv["text"]
+
+        conversation = dict(
+            type="conversation", language="Multilingual", dataset="ShareGPT", conversations=conversations[idx]
+        )
+        conversation_dataset.append(conversation)
+
+    return conversation_dataset
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--dataset",
+        type=str,
+        default="All",
+        choices=["Alpaca", "ShareGPT", "All"],
+        help="which dataset to convert, All will combine Alpaca and ShareGPT",
+    )
+    parser.add_argument("--save_path", type=str, default="dataset.json", help="path to save the converted dataset")
+    args = parser.parse_args()
+
+    conversation_dataset = []
+
+    if args.dataset == "Alpaca":
+        conversation_dataset.extend(generate_alpaca())
+    elif args.dataset == "ShareGPT":
+        conversation_dataset.extend(generate_sharegpt())
+    else:
+        conversation_dataset.extend(generate_alpaca())
+        conversation_dataset.extend(generate_sharegpt())
+
+    for idx, sample in enumerate(conversation_dataset):
+        sample["id"] = idx + 1
+
+    with open(args.save_path, mode="w") as f:
+        json.dump(conversation_dataset, f, indent=4, default=str, ensure_ascii=False)

applications/Chat/examples/generate_prompt_dataset.py

+import argparse
+import json
+import random
+
+random.seed(42)
+
+
+def sample(args):
+    with open(args.dataset_path, mode="r") as f:
+        dataset_list = json.load(f)
+
+    sampled_dataset = [
+        {"instruction": sample["instruction"], "id": idx}
+        for idx, sample in enumerate(random.sample(dataset_list, args.sample_size))
+    ]
+
+    with open(args.save_path, mode="w") as f:
+        json.dump(sampled_dataset, f, indent=4, default=str, ensure_ascii=False)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dataset_path", type=str, default=None, required=True, help="path to the pretrain dataset")
+    parser.add_argument("--save_path", type=str, default="prompt.json", help="path to save the prompt dataset")
+    parser.add_argument("--sample_size", type=int, default=16384, help="size of the prompt dataset")
+    args = parser.parse_args()
+    sample(args)

applications/Chat/examples/inference.py

@@ -2,63 +2,72 @@ import argparse
 
 import torch
 from coati.models.bloom import BLOOMActor
+from coati.models.generation import generate
 from coati.models.gpt import GPTActor
+from coati.models.llama import LlamaActor
 from coati.models.opt import OPTActor
-from coati.models.roberta import RoBERTaActor
-from transformers import AutoTokenizer, RobertaTokenizer
-from transformers.models.gpt2.tokenization_gpt2 import GPT2Tokenizer
+from transformers import AutoTokenizer, BloomTokenizerFast, GPT2Tokenizer, LlamaTokenizer
 
 
 def eval(args):
     # configure model
-    if args.model == 'gpt2':
-        actor = GPTActor(pretrained=args.pretrain).to(torch.cuda.current_device())
-    elif args.model == 'bloom':
-        actor = BLOOMActor(pretrained=args.pretrain).to(torch.cuda.current_device())
-    elif args.model == 'opt':
-        actor = OPTActor(pretrained=args.pretrain).to(torch.cuda.current_device())
-    elif args.model == 'roberta':
-        actor = RoBERTaActor(pretrained=args.pretrain).to(torch.cuda.current_device())
+    if args.model == "gpt2":
+        actor = GPTActor(pretrained=args.pretrain)
+    elif args.model == "bloom":
+        actor = BLOOMActor(pretrained=args.pretrain)
+    elif args.model == "opt":
+        actor = OPTActor(pretrained=args.pretrain)
+    elif args.model == "llama":
+        actor = LlamaActor(pretrained=args.pretrain)
     else:
         raise ValueError(f'Unsupported model "{args.model}"')
 
-    state_dict = torch.load(args.model_path)
-    actor.model.load_state_dict(state_dict)
+    actor.to(torch.cuda.current_device())
+    if args.model_path is not None:
+        state_dict = torch.load(args.model_path)
+        actor.load_state_dict(state_dict)
 
     # configure tokenizer
-    if args.model == 'gpt2':
-        tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
+    if args.model == "gpt2":
+        tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
         tokenizer.pad_token = tokenizer.eos_token
-    elif args.model == 'bloom':
-        tokenizer = AutoTokenizer.from_pretrained('bigscience/bloom-560m')
+    elif args.model == "bloom":
+        tokenizer = BloomTokenizerFast.from_pretrained("bigscience/bloom-560m")
         tokenizer.pad_token = tokenizer.eos_token
-    elif args.model == 'opt':
-        tokenizer = AutoTokenizer.from_pretrained('facebook/opt-350m')
-    elif args.model == 'roberta':
-        tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
+    elif args.model == "opt":
+        tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m")
+        tokenizer.pad_token = tokenizer.eos_token
+    elif args.model == "llama":
+        tokenizer = LlamaTokenizer.from_pretrained("hf-internal-testing/llama-tokenizer")
+        tokenizer.eos_token = "<\s>"
+        tokenizer.pad_token = tokenizer.unk_token
     else:
         raise ValueError(f'Unsupported model "{args.model}"')
 
     actor.eval()
-    input = args.input
-    input_ids = tokenizer.encode(input, return_tensors='pt').to(torch.cuda.current_device())
-    outputs = actor.generate(input_ids,
-                             max_length=args.max_length,
-                             do_sample=True,
-                             top_k=50,
-                             top_p=0.95,
-                             num_return_sequences=1)
+    tokenizer.padding_side = "left"
+    input_ids = tokenizer.encode(args.input, return_tensors="pt").to(torch.cuda.current_device())
+    outputs = generate(
+        actor,
+        input_ids,
+        tokenizer=tokenizer,
+        max_length=args.max_length,
+        do_sample=True,
+        top_k=50,
+        top_p=0.95,
+        num_return_sequences=1,
+    )
     output = tokenizer.batch_decode(outputs[0], skip_special_tokens=True)
-    print(output)
+    print(f"[Output]: {''.join(output)}")
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    parser.add_argument('--model', default='gpt2', choices=['gpt2', 'bloom', 'opt', 'roberta'])
+    parser.add_argument("--model", default="gpt2", choices=["gpt2", "bloom", "opt", "llama"])
     # We suggest to use the pretrained model from HuggingFace, use pretrain to configure model
-    parser.add_argument('--pretrain', type=str, default=None)
-    parser.add_argument('--model_path', type=str, default=None)
-    parser.add_argument('--input', type=str, default='Question: How are you ? Answer:')
-    parser.add_argument('--max_length', type=int, default=100)
+    parser.add_argument("--pretrain", type=str, default=None)
+    parser.add_argument("--model_path", type=str, default=None)
+    parser.add_argument("--input", type=str, default="Question: How are you ? Answer:")
+    parser.add_argument("--max_length", type=int, default=100)
     args = parser.parse_args()
     eval(args)

applications/Chat/examples/ray/1mmt_prompt.py

+import argparse
+import os
+import socket
+from functools import partial
+
+import pandas as pd
+import ray
+from coati.quant import llama_load_quant, low_resource_init
+from coati.ray.detached_trainer_ppo import DetachedPPOTrainer
+from coati.ray.experience_maker_holder import ExperienceMakerHolder
+from coati.ray.utils import (
+    get_actor_from_args,
+    get_critic_from_args,
+    get_reward_model_from_args,
+    get_strategy_from_args,
+    get_tokenizer_from_args,
+)
+from torch.utils.data import DataLoader
+from transformers import AutoConfig
+from transformers.modeling_utils import no_init_weights
+
+
+def get_free_port():
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("", 0))
+        return s.getsockname()[1]
+
+
+def get_local_ip():
+    with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
+        s.connect(("8.8.8.8", 80))
+        return s.getsockname()[0]
+
+
+def main(args):
+    master_addr = str(get_local_ip())
+    # trainer_env_info
+    trainer_port = str(get_free_port())
+    env_info_trainers = [
+        {
+            "local_rank": "0",
+            "rank": str(rank),
+            "world_size": str(args.num_trainers),
+            "master_port": trainer_port,
+            "master_addr": master_addr,
+        }
+        for rank in range(args.num_trainers)
+    ]
+
+    # maker_env_info
+    maker_port = str(get_free_port())
+    env_info_maker = {
+        "local_rank": "0",
+        "rank": "0",
+        "world_size": "1",
+        "master_port": maker_port,
+        "master_addr": master_addr,
+    }
+
+    # configure tokenizer
+    tokenizer = get_tokenizer_from_args(args.model)
+
+    def trainer_model_fn():
+        actor = get_actor_from_args(args.model, args.pretrain).half().cuda()
+        critic = get_critic_from_args(args.model, args.critic_pretrain).half().cuda()
+        return actor, critic
+
+    # configure Trainer
+    trainer_refs = [
+        DetachedPPOTrainer.options(name=f"trainer{i}", num_gpus=1, max_concurrency=2).remote(
+            experience_maker_holder_name_list=["maker1"],
+            strategy_fn=partial(get_strategy_from_args, args.trainer_strategy),
+            model_fn=trainer_model_fn,
+            env_info=env_info_trainer,
+            train_batch_size=args.train_batch_size,
+            buffer_limit=16,
+            eval_performance=True,
+            debug=args.debug,
+            update_lora_weights=not (args.lora_rank == 0),
+        )
+        for i, env_info_trainer in enumerate(env_info_trainers)
+    ]
+
+    def model_fn():
+        actor = get_actor_from_args(args.model, args.pretrain).requires_grad_(False).half().cuda()
+        critic = get_critic_from_args(args.model, args.critic_pretrain).requires_grad_(False).half().cuda()
+        reward_model = get_reward_model_from_args(args.model, args.critic_pretrain).requires_grad_(False).half().cuda()
+        if args.initial_model_quant_ckpt is not None and args.model == "llama":
+            # quantize initial model
+            actor_cfg = AutoConfig.from_pretrained(args.pretrain)
+            with low_resource_init(), no_init_weights():
+                initial_model = get_actor_from_args(args.model, config=actor_cfg)
+            initial_model.model = (
+                llama_load_quant(
+                    initial_model.model, args.initial_model_quant_ckpt, args.quant_bits, args.quant_group_size
+                )
+                .cuda()
+                .requires_grad_(False)
+            )
+        else:
+            initial_model = get_actor_from_args(args.model, args.pretrain).requires_grad_(False).half().cuda()
+        return actor, critic, reward_model, initial_model
+
+    # configure Experience Maker
+    experience_holder_ref = ExperienceMakerHolder.options(name="maker1", num_gpus=1, max_concurrency=2).remote(
+        detached_trainer_name_list=[f"trainer{i}" for i in range(args.num_trainers)],
+        strategy_fn=partial(get_strategy_from_args, args.maker_strategy),
+        model_fn=model_fn,
+        env_info=env_info_maker,
+        experience_batch_size=args.experience_batch_size,
+        kl_coef=0.1,
+        debug=args.debug,
+        update_lora_weights=not (args.lora_rank == 0),
+        # sync_models_from_trainers=True,
+        # generation kwargs:
+        max_length=512,
+        do_sample=True,
+        temperature=1.0,
+        top_k=50,
+        pad_token_id=tokenizer.pad_token_id,
+        eos_token_id=tokenizer.eos_token_id,
+        eval_performance=True,
+        use_cache=True,
+    )
+
+    # uncomment this function if sync_models_from_trainers is True
+    # ray.get([
+    #     trainer_ref.sync_models_to_remote_makers.remote()
+    #     for trainer_ref in trainer_refs
+    # ])
+
+    wait_tasks = []
+
+    total_steps = args.experience_batch_size * args.experience_steps // (args.num_trainers * args.train_batch_size)
+    for trainer_ref in trainer_refs:
+        wait_tasks.append(trainer_ref.fit.remote(total_steps, args.update_steps, args.train_epochs))
+
+    dataset_size = args.experience_batch_size * 4
+
+    def build_dataloader():
+        def tokenize_fn(texts):
+            batch = tokenizer(texts, return_tensors="pt", max_length=96, padding="max_length", truncation=True)
+            return {k: v.cuda() for k, v in batch.items()}
+
+        dataset = pd.read_csv(args.prompt_path)["prompt"]
+        dataloader = DataLoader(dataset=dataset, batch_size=dataset_size, shuffle=True, collate_fn=tokenize_fn)
+        return dataloader
+
+    wait_tasks.append(experience_holder_ref.workingloop.remote(build_dataloader, num_steps=args.experience_steps))
+
+    ray.get(wait_tasks)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--prompt_path", type=str, default=None)
+    parser.add_argument("--num_trainers", type=int, default=1)
+    parser.add_argument(
+        "--trainer_strategy",
+        choices=["ddp", "colossalai_gemini", "colossalai_zero2", "colossalai_gemini_cpu", "colossalai_zero2_cpu"],
+        default="ddp",
+    )
+    parser.add_argument("--maker_strategy", choices=["naive"], default="naive")
+    parser.add_argument("--model", default="gpt2", choices=["gpt2", "bloom", "opt", "llama"])
+    parser.add_argument("--critic_model", default="gpt2", choices=["gpt2", "bloom", "opt", "llama"])
+    parser.add_argument("--pretrain", type=str, default=None)
+    parser.add_argument("--critic_pretrain", type=str, default=None)
+    parser.add_argument("--experience_steps", type=int, default=4)
+    parser.add_argument("--experience_batch_size", type=int, default=8)
+    parser.add_argument("--train_epochs", type=int, default=1)
+    parser.add_argument("--update_steps", type=int, default=2)
+    parser.add_argument("--train_batch_size", type=int, default=8)
+    parser.add_argument("--lora_rank", type=int, default=0, help="low-rank adaptation matrices rank")
+
+    parser.add_argument("--initial_model_quant_ckpt", type=str, default=None)
+    parser.add_argument("--quant_bits", type=int, default=4)
+    parser.add_argument("--quant_group_size", type=int, default=128)
+    parser.add_argument("--debug", action="store_true")
+    args = parser.parse_args()
+    ray.init(namespace=os.environ["RAY_NAMESPACE"], runtime_env={"env_vars": dict(os.environ)})
+    main(args)

applications/Chat/examples/ray/mmmt_prompt.py

+import argparse
+import os
+import socket
+from functools import partial
+
+import pandas as pd
+import ray
+from coati.quant import llama_load_quant, low_resource_init
+from coati.ray.detached_trainer_ppo import DetachedPPOTrainer
+from coati.ray.experience_maker_holder import ExperienceMakerHolder
+from coati.ray.utils import (
+    get_actor_from_args,
+    get_critic_from_args,
+    get_receivers_per_sender,
+    get_reward_model_from_args,
+    get_strategy_from_args,
+)
+from torch.utils.data import DataLoader
+from transformers import AutoConfig, AutoTokenizer
+from transformers.modeling_utils import no_init_weights
+
+
+def get_free_port():
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("", 0))
+        return s.getsockname()[1]
+
+
+def get_local_ip():
+    with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
+        s.connect(("8.8.8.8", 80))
+        return s.getsockname()[0]
+
+
+def main(args):
+    master_addr = str(get_local_ip())
+    # trainer_env_info
+    trainer_port = str(get_free_port())
+    env_info_trainers = [
+        {
+            "local_rank": "0",
+            "rank": str(rank),
+            "world_size": str(args.num_trainers),
+            "master_port": trainer_port,
+            "master_addr": master_addr,
+        }
+        for rank in range(args.num_trainers)
+    ]
+
+    # maker_env_info
+    maker_port = str(get_free_port())
+    env_info_makers = [
+        {
+            "local_rank": "0",
+            "rank": str(rank),
+            "world_size": str(args.num_makers),
+            "master_port": maker_port,
+            "master_addr": master_addr,
+        }
+        for rank in range(args.num_makers)
+    ]
+
+    # configure tokenizer
+    tokenizer = AutoTokenizer.from_pretrained(args.pretrain)
+    tokenizer.pad_token = tokenizer.eos_token
+
+    def model_fn():
+        actor = get_actor_from_args(args.model, args.pretrain).requires_grad_(False).half().cuda()
+        critic = get_critic_from_args(args.model, args.critic_pretrain).requires_grad_(False).half().cuda()
+        reward_model = get_reward_model_from_args(args.model, args.critic_pretrain).requires_grad_(False).half().cuda()
+        if args.initial_model_quant_ckpt is not None and args.model == "llama":
+            # quantize initial model
+            actor_cfg = AutoConfig.from_pretrained(args.pretrain)
+            with low_resource_init(), no_init_weights():
+                initial_model = get_actor_from_args(args.model, config=actor_cfg)
+            initial_model.model = (
+                llama_load_quant(
+                    initial_model.model, args.initial_model_quant_ckpt, args.quant_bits, args.quant_group_size
+                )
+                .cuda()
+                .requires_grad_(False)
+            )
+        else:
+            initial_model = get_actor_from_args(args.model, args.pretrain).requires_grad_(False).half().cuda()
+        return actor, critic, reward_model, initial_model
+
+    # configure Experience Maker
+    experience_holder_refs = [
+        ExperienceMakerHolder.options(name=f"maker{i}", num_gpus=1, max_concurrency=2).remote(
+            detached_trainer_name_list=[
+                f"trainer{x}"
+                for x in get_receivers_per_sender(i, args.num_makers, args.num_trainers, allow_idle_sender=False)
+            ],
+            strategy_fn=partial(get_strategy_from_args, args.maker_strategy),
+            model_fn=model_fn,
+            env_info=env_info_maker,
+            kl_coef=0.1,
+            debug=args.debug,
+            update_lora_weights=not (args.lora_rank == 0),
+            # sync_models_from_trainers=True,
+            # generation kwargs:
+            max_length=512,
+            do_sample=True,
+            temperature=1.0,
+            top_k=50,
+            pad_token_id=tokenizer.pad_token_id,
+            eos_token_id=tokenizer.eos_token_id,
+            eval_performance=True,
+            use_cache=True,
+        )
+        for i, env_info_maker in enumerate(env_info_makers)
+    ]
+
+    def trainer_model_fn():
+        actor = get_actor_from_args(args.model, args.pretrain, lora_rank=args.lora_rank).half().cuda()
+        critic = get_critic_from_args(args.model, args.critic_pretrain, lora_rank=args.lora_rank).half().cuda()
+        return actor, critic
+
+    # configure Trainer
+    trainer_refs = [
+        DetachedPPOTrainer.options(name=f"trainer{i}", num_gpus=1, max_concurrency=2).remote(
+            experience_maker_holder_name_list=[
+                f"maker{x}"
+                for x in get_receivers_per_sender(i, args.num_trainers, args.num_makers, allow_idle_sender=True)
+            ],
+            strategy_fn=partial(get_strategy_from_args, args.trainer_strategy),
+            model_fn=trainer_model_fn,
+            env_info=env_info_trainer,
+            train_batch_size=args.train_batch_size,
+            buffer_limit=16,
+            eval_performance=True,
+            debug=args.debug,
+            update_lora_weights=not (args.lora_rank == 0),
+        )
+        for i, env_info_trainer in enumerate(env_info_trainers)
+    ]
+
+    dataset_size = args.experience_batch_size * 4
+
+    def build_dataloader():
+        def tokenize_fn(texts):
+            batch = tokenizer(texts, return_tensors="pt", max_length=96, padding="max_length", truncation=True)
+            return {k: v.cuda() for k, v in batch.items()}
+
+        dataset = pd.read_csv(args.prompt_path)["prompt"]
+        dataloader = DataLoader(dataset=dataset, batch_size=dataset_size, shuffle=True, collate_fn=tokenize_fn)
+        return dataloader
+
+    # uncomment this function if sync_models_from_trainers is True
+    # ray.get([
+    #     trainer_ref.sync_models_to_remote_makers.remote()
+    #     for trainer_ref in trainer_refs
+    # ])
+
+    wait_tasks = []
+
+    for experience_holder_ref in experience_holder_refs:
+        wait_tasks.append(experience_holder_ref.workingloop.remote(build_dataloader, num_steps=args.experience_steps))
+
+    total_steps = (
+        args.experience_batch_size
+        * args.experience_steps
+        * args.num_makers
+        // (args.num_trainers * args.train_batch_size)
+    )
+    for trainer_ref in trainer_refs:
+        wait_tasks.append(trainer_ref.fit.remote(total_steps, args.update_steps, args.train_epochs))
+
+    ray.get(wait_tasks)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--prompt_path", type=str, default=None)
+    parser.add_argument("--num_makers", type=int, default=1)
+    parser.add_argument("--num_trainers", type=int, default=1)
+    parser.add_argument(
+        "--trainer_strategy",
+        choices=["ddp", "colossalai_gemini", "colossalai_zero2", "colossalai_gemini_cpu", "colossalai_zero2_cpu"],
+        default="ddp",
+    )
+    parser.add_argument("--maker_strategy", choices=["naive"], default="naive")
+    parser.add_argument("--model", default="gpt2", choices=["gpt2", "bloom", "opt", "llama"])
+    parser.add_argument("--critic_model", default="gpt2", choices=["gpt2", "bloom", "opt", "llama"])
+    parser.add_argument("--pretrain", type=str, default=None)
+    parser.add_argument("--critic_pretrain", type=str, default=None)
+    parser.add_argument("--experience_steps", type=int, default=4)
+    parser.add_argument("--experience_batch_size", type=int, default=8)
+    parser.add_argument("--train_epochs", type=int, default=1)
+    parser.add_argument("--update_steps", type=int, default=2)
+    parser.add_argument("--train_batch_size", type=int, default=8)
+    parser.add_argument("--lora_rank", type=int, default=0, help="low-rank adaptation matrices rank")
+
+    parser.add_argument("--initial_model_quant_ckpt", type=str, default=None)
+    parser.add_argument("--quant_bits", type=int, default=4)
+    parser.add_argument("--quant_group_size", type=int, default=128)
+    parser.add_argument("--debug", action="store_true")
+    args = parser.parse_args()
+
+    ray.init(namespace=os.environ["RAY_NAMESPACE"], runtime_env={"env_vars": dict(os.environ)})
+    main(args)

applications/Chat/examples/ray/requirements.txt

+ray

applications/Chat/examples/ray/test_ci.sh

+#!/bin/bash
+
+set -xe
+BASE=$(realpath $(dirname $0))
+
+export RAY_NAMESPACE=admin
+export DATA=/data/scratch/chatgpt/prompts.csv
+
+# install requirements
+pip install -r ${BASE}/requirements.txt
+
+python ${BASE}/mmmt_prompt.py --prompt_path $DATA --num_makers 2 --num_trainers 2 --trainer_strategy colossalai_gemini --model opt --critic_model opt --pretrain facebook/opt-350m --critic_pretrain facebook/opt-125m --experience_batch_size 4 --train_batch_size 2

applications/Chat/examples/requirements.txt

 pandas>=1.4.1
 sentencepiece
+colossalai==0.3.3

applications/Chat/examples/test_ci.sh

-#!/usr/bin/env bash
-
-set -xue
-
-if [ -z "$SFT_DATASET" ]; then
-    echo "Please set \$SFT_DATASET to the path to sft dataset."
-    exit 1
-fi
-
-if [ -z "$PROMPT_PATH" ]; then
-    echo "Please set \$PROMPT_PATH to the path to prompts csv."
-    exit 1
-fi
-
-if [ -z "$PRETRAIN_DATASET" ]; then
-    echo "Please set \$PRETRAIN_DATASET to the path to alpaca data."
-    exit 1
-fi
-
-BASE=$(realpath $(dirname $0))
-
-export OMP_NUM_THREADS=8
-
-# install requirements
-pip install -r ${BASE}/requirements.txt
-
-wandb init -m offline
-
-# train sft
-torchrun --standalone --nproc_per_node=4 ${BASE}/train_sft.py --pretrain 'bigscience/bloom-560m' \
-        --model 'bloom' --strategy colossalai_zero2 --lora_rank 4\
-        --dataset $SFT_DATASET --max_datasets_size 512 --max_epochs 1 \
-        --save_path ${BASE}/output
-rm -rf ${BASE}/output
-
-torchrun --standalone --nproc_per_node=4 ${BASE}/train_sft.py --pretrain 'gpt2' \
-        --model 'gpt2' --strategy colossalai_zero2 \
-        --dataset $SFT_DATASET --max_datasets_size 512 --max_epochs 1 \
-        --save_path ${BASE}/output
-rm -rf ${BASE}/output
-
-torchrun --standalone --nproc_per_node=4 ${BASE}/train_sft.py --pretrain 'facebook/opt-350m' \
-        --model 'opt' --strategy colossalai_zero2 --lora_rank 4\
-        --dataset $SFT_DATASET --max_datasets_size 512 --max_epochs 1 \
-        --save_path ${BASE}/output
-rm -rf ${BASE}/output
-
-torchrun --standalone --nproc_per_node=4 ${BASE}/train_sft.py --pretrain 'gpt2' \
-        --model 'gpt2' --strategy ddp --lora_rank 4\
-        --dataset $SFT_DATASET --max_datasets_size 512 --max_epochs 1 \
-        --save_path ${BASE}/output
-
-#torchrun --standalone --nproc_per_node=4 ${BASE}/train_sft.py --pretrain 'facebook/opt-350m' \
-#        --model 'opt' --strategy naive \
-#        --dataset $SFT_DATASET --max_datasets_size 512 --max_epochs 1 \
-#        --save_path ${BASE}/output
-
-rm -rf ${BASE}/output
-
-# train rm
-torchrun --standalone --nproc_per_node=2 ${BASE}/train_reward_model.py \
-                            --pretrain 'facebook/opt-350m' --model 'opt' \
-                            --strategy colossalai_zero2 --loss_fn 'log_sig'\
-                            --dataset 'Anthropic/hh-rlhf' --subset 'harmless-base' \
-                            --test True --lora_rank 0 \
-                            --save_path ${BASE}/rm_ckpt_opt.pt
-
-torchrun --standalone --nproc_per_node=2 ${BASE}/train_reward_model.py \
-                            --pretrain 'gpt2' --model 'gpt2' \
-                            --strategy colossalai_zero2 --loss_fn 'log_exp' \
-                            --dataset 'Dahoas/rm-static' \
-                            --test True  --lora_rank 0 \
-                            --save_path ${BASE}/rm_ckpt_gpt.pt
-
-torchrun --standalone --nproc_per_node=2 ${BASE}/train_reward_model.py \
-                            --pretrain 'gpt2' --model 'gpt2' \
-                            --strategy ddp --loss_fn 'log_exp' \
-                            --dataset 'Dahoas/rm-static' \
-                            --test True --lora_rank 4 \
-                            --save_path ${BASE}/rm_ckpt.pt
-rm -rf ${BASE}/rm_ckpt.pt
-
-torchrun --standalone --nproc_per_node=2 ${BASE}/train_reward_model.py \
-                            --pretrain 'bigscience/bloom-560m' --model 'bloom' \
-                            --strategy colossalai_zero2 --loss_fn 'log_sig' \
-                            --dataset 'Anthropic/hh-rlhf' --subset 'harmless-base' \
-                            --test True --lora_rank 4 \
-                            --save_path ${BASE}/rm_ckpt.pt
-rm -rf ${BASE}/rm_ckpt.pt
-
-torchrun --standalone --nproc_per_node=2 ${BASE}/train_reward_model.py \
-                            --pretrain 'microsoft/deberta-v3-large' --model 'deberta' \
-                            --strategy colossalai_zero2 --loss_fn 'log_sig' \
-                            --dataset 'Anthropic/hh-rlhf' --subset 'harmless-base' \
-                            --test True --lora_rank 4 \
-                            --save_path ${BASE}/rm_ckpt.pt
-rm -rf ${BASE}/rm_ckpt.pt
-
-torchrun --standalone --nproc_per_node=2 ${BASE}/train_reward_model.py \
-                            --pretrain 'roberta-base' --model 'roberta' \
-                            --strategy colossalai_zero2 --loss_fn 'log_exp'\
-                            --dataset 'Anthropic/hh-rlhf' --subset 'harmless-base'\
-                            --test True --lora_rank 4 \
-                            --save_path ${BASE}/rm_ckpt.pt
-
-rm -rf ${BASE}/rm_ckpt.pt
-
-torchrun --standalone --nproc_per_node=2 ${BASE}/train_prompts.py --prompt_dataset $PROMPT_PATH --pretrain_dataset $PRETRAIN_DATASET \
-        --strategy colossalai_zero2 --num_episodes 1 --max_timesteps 2 \
-        --update_timesteps 2 --max_epochs 1 --train_batch_size 2 \
-        --pretrain 'facebook/opt-350m' --model opt \
-        --rm_pretrain 'facebook/opt-350m' \
-        --rm_path ${BASE}/rm_ckpt_opt.pt \
-        --save_path ${BASE}/actor_checkpoint_prompts.pt
-rm -rf ${BASE}/rm_ckpt_opt.pt
-
-torchrun --standalone --nproc_per_node=2 ${BASE}/train_prompts.py --prompt_dataset $PROMPT_PATH --pretrain_dataset $PRETRAIN_DATASET \
-         --strategy colossalai_zero2 --num_episodes 1 --max_timesteps 2 \
-         --update_timesteps 2 --max_epochs 1 --train_batch_size 2 \
-         --pretrain 'gpt2' --model gpt2 \
-         --rm_pretrain 'gpt2' \
-         --rm_path ${BASE}/rm_ckpt_gpt.pt \
-         --save_path ${BASE}/actor_checkpoint_prompts.pt
-rm -rf ${BASE}/rm_ckpt_gpt.pt
-
-rm -rf ${BASE}/actor_checkpoint_prompts.pt

applications/Chat/examples/train_prompts.py

 import argparse
+import warnings
 
-import pandas as pd
 import torch
 import torch.distributed as dist
-from coati.dataset import DataCollatorForSupervisedDataset, PromptDataset, SupervisedDataset
+from coati.dataset import PromptDataset, SupervisedDataset
 from coati.models.bloom import BLOOMRM, BLOOMActor, BLOOMCritic
 from coati.models.gpt import GPTRM, GPTActor, GPTCritic
 from coati.models.llama import LlamaActor, LlamaCritic, LlamaRM
 from coati.models.opt import OPTRM, OPTActor, OPTCritic
-from coati.models.roberta import RoBERTaActor, RoBERTaCritic, RoBERTaRM
 from coati.trainer import PPOTrainer
-from coati.trainer.strategies import ColossalAIStrategy, DDPStrategy, NaiveStrategy
-from coati.utils import prepare_llama_tokenizer_and_embedding
+from coati.trainer.strategies import DDPStrategy, GeminiStrategy, LowLevelZeroStrategy
 from torch.optim import Adam
 from torch.utils.data import DataLoader
 from torch.utils.data.distributed import DistributedSampler
-from transformers import AutoTokenizer, BloomTokenizerFast, GPT2Tokenizer, LlamaTokenizer, RobertaTokenizer
+from transformers import AutoTokenizer, BloomTokenizerFast, GPT2Tokenizer, LlamaTokenizer
 
 from colossalai.nn.optimizer import HybridAdam
 
 
 def main(args):
     # configure strategy
-    if args.strategy == 'naive':
-        strategy = NaiveStrategy()
-    elif args.strategy == 'ddp':
+    if args.strategy == "ddp":
         strategy = DDPStrategy()
-    elif args.strategy == 'colossalai_gemini':
-        strategy = ColossalAIStrategy(stage=3, placement_policy='cuda', initial_scale=2**5)
-    elif args.strategy == 'colossalai_zero2':
-        strategy = ColossalAIStrategy(stage=2, placement_policy='cuda')
+    elif args.strategy == "colossalai_gemini":
+        strategy = GeminiStrategy(placement_policy="static", initial_scale=2**5)
+    elif args.strategy == "colossalai_zero2":
+        strategy = LowLevelZeroStrategy(stage=2, placement_policy="cuda")
     else:
         raise ValueError(f'Unsupported strategy "{args.strategy}"')
 
     if args.rm_path is not None:
-        state_dict = torch.load(args.rm_path, map_location='cpu')
-
-    # configure model
-    if args.model == 'gpt2':
-        initial_model = GPTActor(pretrained=args.pretrain)
-    elif args.model == 'bloom':
-        initial_model = BLOOMActor(pretrained=args.pretrain)
-    elif args.model == 'opt':
-        initial_model = OPTActor(pretrained=args.pretrain)
-    elif args.model == 'llama':
-        initial_model = LlamaActor(pretrained=args.pretrain)
-    elif args.model == 'roberta':
-        initial_model = RoBERTaActor(pretrained=args.pretrain)
-    else:
-        raise ValueError(f'Unsupported actor model "{args.model}"')
+        warnings.warn("LoRA weights should be merged with the model weights")
+        state_dict = torch.load(args.rm_path, map_location="cpu")
 
-    if args.rm_model == None:
-        rm_model_name = args.model
-    else:
-        rm_model_name = args.rm_model
-
-    if rm_model_name == 'gpt2':
-        reward_model = GPTRM(pretrained=args.rm_pretrain)
-    elif rm_model_name == 'bloom':
-        reward_model = BLOOMRM(pretrained=args.rm_pretrain)
-    elif rm_model_name == 'opt':
-        reward_model = OPTRM(pretrained=args.rm_pretrain)
-    elif rm_model_name == 'llama':
-        reward_model = LlamaRM(pretrained=args.rm_pretrain)
-    elif rm_model_name == 'roberta':
-        reward_model = RoBERTaRM(pretrained=args.rm_pretrain)
-    else:
-        raise ValueError(f'Unsupported reward model "{rm_model_name}"')
+    if args.lora_rank > 0:
+        warnings.warn("Lora is not supported yet.")
+        args.lora_rank = 0
 
-    if args.rm_path is not None:
-        reward_model.load_state_dict(state_dict)
+    with strategy.model_init_context():
+        # configure model
+        if args.model == "gpt2":
+            initial_model = GPTActor(pretrained=args.pretrain)
+        elif args.model == "bloom":
+            initial_model = BLOOMActor(pretrained=args.pretrain)
+        elif args.model == "opt":
+            initial_model = OPTActor(pretrained=args.pretrain)
+        elif args.model == "llama":
+            initial_model = LlamaActor(pretrained=args.pretrain)
+        else:
+            raise ValueError(f'Unsupported actor model "{args.model}"')
+
+        if args.rm_model is None:
+            rm_model_name = args.model
+        else:
+            rm_model_name = args.rm_model
+
+        if rm_model_name == "gpt2":
+            reward_model = GPTRM(pretrained=args.rm_pretrain, lora_rank=args.lora_rank)
+        elif rm_model_name == "bloom":
+            reward_model = BLOOMRM(pretrained=args.rm_pretrain, lora_rank=args.lora_rank)
+        elif rm_model_name == "opt":
+            reward_model = OPTRM(pretrained=args.rm_pretrain, lora_rank=args.lora_rank)
+        elif rm_model_name == "llama":
+            reward_model = LlamaRM(pretrained=args.rm_pretrain, lora_rank=args.lora_rank)
+        else:
+            raise ValueError(f'Unsupported reward model "{rm_model_name}"')
+
+        if args.rm_path is not None:
+            reward_model.load_state_dict(state_dict, strict=False)
 
-    initial_model.to(torch.float16).to(torch.cuda.current_device())
-    reward_model.to(torch.float16).to(torch.cuda.current_device())
+        initial_model.to(torch.bfloat16).to(torch.cuda.current_device())
+        reward_model.to(torch.bfloat16).to(torch.cuda.current_device())
 
-    with strategy.model_init_context():
-        if args.model == 'gpt2':
+        if args.model == "gpt2":
             actor = GPTActor(pretrained=args.pretrain, lora_rank=args.lora_rank)
-        elif args.model == 'bloom':
+        elif args.model == "bloom":
             actor = BLOOMActor(pretrained=args.pretrain, lora_rank=args.lora_rank)
-        elif args.model == 'opt':
+        elif args.model == "opt":
             actor = OPTActor(pretrained=args.pretrain, lora_rank=args.lora_rank)
-        elif args.model == 'llama':
+        elif args.model == "llama":
             actor = LlamaActor(pretrained=args.pretrain, lora_rank=args.lora_rank)
-        elif args.model == 'roberta':
-            actor = RoBERTaActor(pretrained=args.pretrain, lora_rank=args.lora_rank)
         else:
             raise ValueError(f'Unsupported actor model "{args.model}"')
 
-        if rm_model_name == 'gpt2':
-            critic = GPTCritic(pretrained=args.rm_pretrain, lora_rank=args.lora_rank, use_action_mask=True)
-        elif rm_model_name == 'bloom':
-            critic = BLOOMCritic(pretrained=args.rm_pretrain, lora_rank=args.lora_rank, use_action_mask=True)
-        elif rm_model_name == 'opt':
-            critic = OPTCritic(pretrained=args.rm_pretrain, lora_rank=args.lora_rank, use_action_mask=True)
-        elif rm_model_name == 'llama':
-            critic = LlamaCritic(pretrained=args.rm_pretrain, lora_rank=args.lora_rank, use_action_mask=True)
-        elif rm_model_name == 'roberta':
-            critic = RoBERTaCritic(pretrained=args.rm_pretrain, lora_rank=args.lora_rank, use_action_mask=True)
+        if rm_model_name == "gpt2":
+            critic = GPTCritic(pretrained=args.rm_pretrain, lora_rank=args.lora_rank)
+        elif rm_model_name == "bloom":
+            critic = BLOOMCritic(pretrained=args.rm_pretrain, lora_rank=args.lora_rank)
+        elif rm_model_name == "opt":
+            critic = OPTCritic(pretrained=args.rm_pretrain, lora_rank=args.lora_rank)
+        elif rm_model_name == "llama":
+            critic = LlamaCritic(pretrained=args.rm_pretrain, lora_rank=args.lora_rank)
         else:
             raise ValueError(f'Unsupported reward model "{rm_model_name}"')
 
         if args.rm_path is not None:
-            critic.load_state_dict(state_dict)
+            critic.load_state_dict(state_dict, strict=False)
             del state_dict
 
-    if args.strategy != 'colossalai_gemini':
-        critic.to(torch.float16).to(torch.cuda.current_device())
-        actor.to(torch.float16).to(torch.cuda.current_device())
+        actor.to(torch.bfloat16).to(torch.cuda.current_device())
+        critic.to(torch.bfloat16).to(torch.cuda.current_device())
 
     # configure optimizer
-    if args.strategy.startswith('colossalai'):
-        actor_optim = HybridAdam(actor.parameters(), lr=1e-7)
-        critic_optim = HybridAdam(critic.parameters(), lr=1e-7)
+    if args.strategy.startswith("colossalai"):
+        actor_optim = HybridAdam(actor.parameters(), lr=args.lr)
+        critic_optim = HybridAdam(critic.parameters(), lr=args.lr)
     else:
-        actor_optim = Adam(actor.parameters(), lr=1e-7)
-        critic_optim = Adam(critic.parameters(), lr=1e-7)
+        actor_optim = Adam(actor.parameters(), lr=args.lr)
+        critic_optim = Adam(critic.parameters(), lr=args.lr)
 
     # configure tokenizer
-    if args.model == 'gpt2':
-        tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
-    elif args.model == 'bloom':
-        tokenizer = BloomTokenizerFast.from_pretrained('bigscience/bloom-560m')
-    elif args.model == 'opt':
-        tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m")
-    elif args.model == 'llama':
-        tokenizer = LlamaTokenizer.from_pretrained(args.pretrain)
-        tokenizer.eos_token = '<\s>'
-    elif args.model == 'roberta':
-        tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
+    if args.model == "gpt2":
+        tokenizer = GPT2Tokenizer.from_pretrained("gpt2" if args.tokenizer is None else args.tokenizer)
+        tokenizer.pad_token = tokenizer.eos_token
+    elif args.model == "bloom":
+        tokenizer = BloomTokenizerFast.from_pretrained(
+            "bigscience/bloom-560m" if args.tokenizer is None else args.tokenizer
+        )
+        tokenizer.pad_token = tokenizer.eos_token
+    elif args.model == "opt":
+        tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m" if args.tokenizer is None else args.tokenizer)
+        tokenizer.pad_token = tokenizer.eos_token
+    elif args.model == "llama":
+        tokenizer = LlamaTokenizer.from_pretrained(
+            "hf-internal-testing/llama-tokenizer" if args.tokenizer is None else args.tokenizer
+        )
+        tokenizer.eos_token = "<\s>"
+        tokenizer.pad_token = tokenizer.unk_token
     else:
         raise ValueError(f'Unsupported model "{args.model}"')
-
-    if args.model == 'llama':
-        tokenizer = prepare_llama_tokenizer_and_embedding(tokenizer, actor)
-    else:
-        tokenizer.pad_token = tokenizer.eos_token
-
-    data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)
-
-    prompt_dataset = PromptDataset(tokenizer=tokenizer, data_path=args.prompt_dataset, max_datasets_size=16384)
+    # NOTE: generate() requires padding_side to be "left"
+    tokenizer.padding_side = "left"
+
+    prompt_dataset = PromptDataset(
+        tokenizer=tokenizer,
+        data_path=args.prompt_dataset,
+        max_datasets_size=args.max_datasets_size,
+        max_length=args.max_input_len,
+    )
     if dist.is_initialized() and dist.get_world_size() > 1:
         prompt_sampler = DistributedSampler(prompt_dataset, shuffle=True, seed=42, drop_last=True)
     else:
         prompt_sampler = None
-    prompt_dataloader = DataLoader(prompt_dataset,
-                                   shuffle=(prompt_sampler is None),
-                                   sampler=prompt_sampler,
-                                   batch_size=args.experience_batch_size)
-
-    pretrain_dataset = SupervisedDataset(tokenizer=tokenizer,
-                                         data_path=args.pretrain_dataset,
-                                         max_datasets_size=16384,
-                                         max_length=args.max_input_len)
+    prompt_dataloader = DataLoader(
+        prompt_dataset, shuffle=(prompt_sampler is None), sampler=prompt_sampler, batch_size=args.experience_batch_size
+    )
+
+    pretrain_dataset = SupervisedDataset(
+        tokenizer=tokenizer,
+        data_path=args.pretrain_dataset,
+        max_datasets_size=args.max_datasets_size,
+        max_length=args.max_input_len,
+    )
     if dist.is_initialized() and dist.get_world_size() > 1:
         pretrain_sampler = DistributedSampler(pretrain_dataset, shuffle=True, seed=42, drop_last=True)
     else:
         pretrain_sampler = None
-    pretrain_dataloader = DataLoader(pretrain_dataset,
-                                     shuffle=(pretrain_sampler is None),
-                                     sampler=pretrain_sampler,
-                                     batch_size=args.ptx_batch_size,
-                                     collate_fn=data_collator)
+    pretrain_dataloader = DataLoader(
+        pretrain_dataset, shuffle=(pretrain_sampler is None), sampler=pretrain_sampler, batch_size=args.ptx_batch_size
+    )
 
-    (actor, actor_optim), (critic, critic_optim) = strategy.prepare((actor, actor_optim), (critic, critic_optim))
+    # NOTE: For small models like opt-1.3b, reward model and initial model are not required to be parallelized.
+    (actor, actor_optim), (critic, critic_optim), reward_model, initial_model = strategy.prepare(
+        (actor, actor_optim), (critic, critic_optim), reward_model, initial_model
+    )
 
     # configure trainer
     trainer = PPOTrainer(
@@ -174,60 +174,76 @@ def main(args):
         initial_model,
         actor_optim,
         critic_optim,
+        tokenizer=tokenizer,
         kl_coef=args.kl_coef,
         ptx_coef=args.ptx_coef,
-        max_epochs=args.max_epochs,
         train_batch_size=args.train_batch_size,
         max_length=args.max_seq_len,
         use_cache=True,
         do_sample=True,
         temperature=1.0,
         top_k=50,
-        pad_token_id=tokenizer.pad_token_id,
-        eos_token_id=tokenizer.eos_token_id,
+        offload_inference_models=args.strategy != "colossalai_gemini",
     )
 
-    trainer.fit(prompt_dataloader=prompt_dataloader,
-                pretrain_dataloader=pretrain_dataloader,
-                num_episodes=args.num_episodes,
-                max_timesteps=args.max_timesteps,
-                update_timesteps=args.update_timesteps)
+    trainer.fit(
+        num_episodes=args.num_episodes,
+        num_collect_steps=args.num_collect_steps,
+        num_update_steps=args.num_update_steps,
+        prompt_dataloader=prompt_dataloader,
+        pretrain_dataloader=pretrain_dataloader,
+        log_dir=args.log_dir,
+        use_wandb=args.use_wandb,
+    )
+
+    if args.lora_rank > 0 and args.merge_lora_weights:
+        from coati.models.lora import LORA_MANAGER
 
+        # NOTE: set model to eval to merge LoRA weights
+        LORA_MANAGER.merge_weights = True
+        actor.eval()
     # save model checkpoint after fitting
-    strategy.save_model(actor, args.save_path, only_rank0=True)
+    strategy.save_pretrained(actor, path=args.save_path)
     # save optimizer checkpoint on all ranks
     if args.need_optim_ckpt:
-        strategy.save_optimizer(actor_optim,
-                                'actor_optim_checkpoint_prompts_%d.pt' % (torch.cuda.current_device()),
-                                only_rank0=False)
+        strategy.save_optimizer(
+            actor_optim, "actor_optim_checkpoint_prompts_%d.pt" % (torch.cuda.current_device()), only_rank0=False
+        )
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    parser.add_argument('--prompt_dataset', type=str, default=None, help='path to the prompt dataset')
-    parser.add_argument('--pretrain_dataset', type=str, default=None, help='path to the pretrained dataset')
-    parser.add_argument('--strategy',
-                        choices=['naive', 'ddp', 'colossalai_gemini', 'colossalai_zero2'],
-                        default='colossalai_zero2',
-                        help='strategy to use')
-    parser.add_argument('--model', default='gpt2', choices=['gpt2', 'bloom', 'opt', 'llama', 'roberta'])
-    parser.add_argument('--pretrain', type=str, default=None)
-    parser.add_argument('--rm_model', default=None, choices=['gpt2', 'bloom', 'opt', 'llama', 'roberta'])
-    parser.add_argument('--rm_path', type=str, default=None)
-    parser.add_argument('--rm_pretrain', type=str, default=None)
-    parser.add_argument('--save_path', type=str, default='actor_checkpoint_prompts')
-    parser.add_argument('--need_optim_ckpt', type=bool, default=False)
-    parser.add_argument('--num_episodes', type=int, default=10)
-    parser.add_argument('--max_timesteps', type=int, default=10)
-    parser.add_argument('--update_timesteps', type=int, default=10)
-    parser.add_argument('--max_epochs', type=int, default=5)
-    parser.add_argument('--train_batch_size', type=int, default=8)
-    parser.add_argument('--ptx_batch_size', type=int, default=1)
-    parser.add_argument('--experience_batch_size', type=int, default=8)
-    parser.add_argument('--lora_rank', type=int, default=0, help="low-rank adaptation matrices rank")
-    parser.add_argument('--kl_coef', type=float, default=0.1)
-    parser.add_argument('--ptx_coef', type=float, default=0.9)
-    parser.add_argument('--max_input_len', type=int, default=96)
-    parser.add_argument('--max_seq_len', type=int, default=128)
+    parser.add_argument("--prompt_dataset", type=str, default=None, help="path to the prompt dataset")
+    parser.add_argument("--pretrain_dataset", type=str, default=None, help="path to the pretrained dataset")
+    parser.add_argument("--max_datasets_size", type=int, default=50000)
+    parser.add_argument(
+        "--strategy",
+        choices=["ddp", "colossalai_gemini", "colossalai_zero2"],
+        default="colossalai_zero2",
+        help="strategy to use",
+    )
+    parser.add_argument("--model", default="gpt2", choices=["gpt2", "bloom", "opt", "llama"])
+    parser.add_argument("--tokenizer", type=str, default=None)
+    parser.add_argument("--pretrain", type=str, default=None)
+    parser.add_argument("--rm_model", default=None, choices=["gpt2", "bloom", "opt", "llama"])
+    parser.add_argument("--rm_path", type=str, default=None)
+    parser.add_argument("--rm_pretrain", type=str, default=None)
+    parser.add_argument("--save_path", type=str, default="actor_checkpoint_prompts")
+    parser.add_argument("--need_optim_ckpt", type=bool, default=False)
+    parser.add_argument("--num_episodes", type=int, default=10)
+    parser.add_argument("--num_collect_steps", type=int, default=10)
+    parser.add_argument("--num_update_steps", type=int, default=5)
+    parser.add_argument("--train_batch_size", type=int, default=8)
+    parser.add_argument("--ptx_batch_size", type=int, default=1)
+    parser.add_argument("--experience_batch_size", type=int, default=8)
+    parser.add_argument("--lora_rank", type=int, default=0, help="low-rank adaptation matrices rank")
+    parser.add_argument("--merge_lora_weights", type=bool, default=True)
+    parser.add_argument("--lr", type=float, default=1e-7)
+    parser.add_argument("--kl_coef", type=float, default=0.1)
+    parser.add_argument("--ptx_coef", type=float, default=0.9)
+    parser.add_argument("--max_input_len", type=int, default=96)
+    parser.add_argument("--max_seq_len", type=int, default=128)
+    parser.add_argument("--log_dir", default="logs", type=str)
+    parser.add_argument("--use_wandb", default=False, action="store_true")
     args = parser.parse_args()
     main(args)

applications/Chat/examples/train_prompts.sh

 set_n_least_used_CUDA_VISIBLE_DEVICES() {
     local n=${1:-"9999"}
     echo "GPU Memory Usage:"
-    local FIRST_N_GPU_IDS=$(nvidia-smi --query-gpu=memory.used --format=csv \
-        | tail -n +2 \
-        | nl -v 0 \
-        | tee /dev/tty \
-        | sort -g -k 2 \
-        | awk '{print $1}' \
-        | head -n $n)
+    local FIRST_N_GPU_IDS=$(nvidia-smi --query-gpu=memory.used --format=csv |
+        tail -n +2 |
+        nl -v 0 |
+        tee /dev/tty |
+        sort -g -k 2 |
+        awk '{print $1}' |
+        head -n $n)
     export CUDA_VISIBLE_DEVICES=$(echo $FIRST_N_GPU_IDS | sed 's/ /,/g')
     echo "Now CUDA_VISIBLE_DEVICES is set to:"
     echo "CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES"
@@ -17,4 +17,9 @@ set_n_least_used_CUDA_VISIBLE_DEVICES 2
 
 # torchrun --standalone --nproc_per_node=2 train_prompts.py prompts.csv --strategy colossalai_zero2
 
-torchrun --standalone --nproc_per_node=2 train_prompts.py --prompt_dataset /path/to/data.json --strategy colossalai_zero2
+torchrun --standalone --nproc_per_node=2 train_prompts.py \
+    --pretrain_dataset /path/to/data.json \
+    --prompt_dataset /path/to/data.json \
+    --strategy colossalai_zero2 \
+    --num_episodes 1 --num_collect_steps 2 --num_update_steps 1 \
+    --train_batch_size 2

applications/Chat/examples/train_reward_model.py

 import argparse
-from random import randint
+import warnings
 
-import loralib as lora
 import torch
 import torch.distributed as dist
 from coati.dataset import HhRlhfDataset, RmStaticDataset
 from coati.models import LogExpLoss, LogSigLoss
-from coati.models.base import RewardModel
 from coati.models.bloom import BLOOMRM
-from coati.models.deberta import DebertaRM
 from coati.models.gpt import GPTRM
 from coati.models.llama import LlamaRM
 from coati.models.opt import OPTRM
-from coati.models.roberta import RoBERTaRM
 from coati.trainer import RewardModelTrainer
-from coati.trainer.strategies import ColossalAIStrategy, DDPStrategy, NaiveStrategy
-from coati.utils import prepare_llama_tokenizer_and_embedding
+from coati.trainer.strategies import DDPStrategy, GeminiStrategy, LowLevelZeroStrategy
 from datasets import load_dataset
 from torch.optim import Adam
+from torch.optim.lr_scheduler import CosineAnnealingLR
 from torch.utils.data import DataLoader
 from torch.utils.data.distributed import DistributedSampler
-from transformers import AutoTokenizer, BloomTokenizerFast, DebertaV2Tokenizer, LlamaTokenizer, RobertaTokenizer
+from transformers import AutoTokenizer, BloomTokenizerFast, LlamaTokenizer
 from transformers.models.gpt2.tokenization_gpt2 import GPT2Tokenizer
 
 from colossalai.nn.optimizer import HybridAdam
@@ -28,72 +24,69 @@ from colossalai.nn.optimizer import HybridAdam
 
 def train(args):
     # configure strategy
-    if args.strategy == 'naive':
-        strategy = NaiveStrategy()
-    elif args.strategy == 'ddp':
+    if args.strategy == "ddp":
         strategy = DDPStrategy()
-    elif args.strategy == 'colossalai_gemini':
-        strategy = ColossalAIStrategy(stage=3, placement_policy='cuda')
-    elif args.strategy == 'colossalai_zero2':
-        strategy = ColossalAIStrategy(stage=2, placement_policy='cuda')
+    elif args.strategy == "colossalai_gemini":
+        strategy = GeminiStrategy(placement_policy="auto")
+    elif args.strategy == "colossalai_zero2":
+        strategy = LowLevelZeroStrategy(stage=2, placement_policy="cuda")
     else:
         raise ValueError(f'Unsupported strategy "{args.strategy}"')
 
     # configure model
+    if args.lora_rank > 0:
+        warnings.warn("Lora is not supported yet.")
+        args.lora_rank = 0
+
     with strategy.model_init_context():
-        if args.model == 'bloom':
-            model = BLOOMRM(pretrained=args.pretrain, lora_rank=args.lora_rank).to(torch.cuda.current_device())
-        elif args.model == 'opt':
-            model = OPTRM(pretrained=args.pretrain, lora_rank=args.lora_rank).to(torch.cuda.current_device())
-        elif args.model == 'gpt2':
-            model = GPTRM(pretrained=args.pretrain, lora_rank=args.lora_rank).to(torch.cuda.current_device())
-        elif args.model == 'deberta':
-            model = DebertaRM(pretrained=args.pretrain, lora_rank=args.lora_rank).to(torch.cuda.current_device())
-        elif args.model == 'llama':
-            model = LlamaRM(pretrained=args.pretrain, lora_rank=args.lora_rank).to(torch.cuda.current_device())
-        elif args.model == 'roberta':
-            model = RoBERTaRM(pretrained=args.pretrain, lora_rank=args.lora_rank).to(torch.cuda.current_device())
+        if args.model == "bloom":
+            model = BLOOMRM(pretrained=args.pretrain, lora_rank=args.lora_rank)
+        elif args.model == "opt":
+            model = OPTRM(pretrained=args.pretrain, lora_rank=args.lora_rank)
+        elif args.model == "gpt2":
+            model = GPTRM(pretrained=args.pretrain, lora_rank=args.lora_rank)
+        elif args.model == "llama":
+            model = LlamaRM(pretrained=args.pretrain, lora_rank=args.lora_rank)
         else:
             raise ValueError(f'Unsupported model "{args.model}"')
 
+        model.to(torch.bfloat16).to(torch.cuda.current_device())
+
         if args.model_path is not None:
             state_dict = torch.load(args.model_path)
             model.load_state_dict(state_dict)
 
-    model = model.to(torch.float16)
-
     # configure tokenizer
-    if args.model == 'gpt2':
-        tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
-    elif args.model == 'bloom':
-        tokenizer = BloomTokenizerFast.from_pretrained('bigscience/bloom-560m')
-    elif args.model == 'opt':
-        tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m")
-    elif args.model == 'deberta':
-        tokenizer = DebertaV2Tokenizer.from_pretrained('microsoft/deberta-v3-large')
-    elif args.model == 'llama':
-        tokenizer = LlamaTokenizer.from_pretrained(args.pretrain)
-    elif args.model == 'roberta':
-        tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
+    if args.model == "gpt2":
+        tokenizer = GPT2Tokenizer.from_pretrained("gpt2" if args.tokenizer is None else args.tokenizer)
+        tokenizer.pad_token = tokenizer.eos_token
+    elif args.model == "bloom":
+        tokenizer = BloomTokenizerFast.from_pretrained(
+            "bigscience/bloom-560m" if args.tokenizer is None else args.tokenizer
+        )
+        tokenizer.pad_token = tokenizer.eos_token
+    elif args.model == "opt":
+        tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m" if args.tokenizer is None else args.tokenizer)
+        tokenizer.pad_token = tokenizer.eos_token
+    elif args.model == "llama":
+        tokenizer = LlamaTokenizer.from_pretrained(
+            "hf-internal-testing/llama-tokenizer" if args.tokenizer is None else args.tokenizer
+        )
+        tokenizer.eos_token = "<\s>"
+        tokenizer.pad_token = tokenizer.unk_token
     else:
         raise ValueError(f'Unsupported model "{args.model}"')
-    max_len = args.max_len
-
-    if args.model == 'llama':
-        tokenizer = prepare_llama_tokenizer_and_embedding(tokenizer, model)
-    else:
-        tokenizer.pad_token = tokenizer.eos_token
 
     # configure optimizer
-    if args.strategy.startswith('colossalai'):
-        optim = HybridAdam(model.parameters(), lr=5e-6)
+    if args.strategy.startswith("colossalai"):
+        optim = HybridAdam(model.parameters(), lr=args.lr)
     else:
-        optim = Adam(model.parameters(), lr=5e-6)
+        optim = Adam(model.parameters(), lr=args.lr)
 
     # configure loss function
-    if args.loss_fn == 'log_sig':
+    if args.loss_fn == "log_sig":
         loss_fn = LogSigLoss()
-    elif args.loss_fn == 'log_exp':
+    elif args.loss_fn == "log_exp":
         loss_fn = LogExpLoss()
     else:
         raise ValueError(f'Unsupported loss function "{args.loss_fn}"')
@@ -104,107 +97,112 @@ def train(args):
     else:
         data = load_dataset(args.dataset)
 
-    if args.test:
-        train_data = data['train'].select(range(100))
-        eval_data = data['test'].select(range(10))
-    else:
-        train_data = data['train']
-        eval_data = data['test']
-    valid_data = data['test'].select((randint(0, len(eval_data) - 1) for _ in range(len(eval_data) // 5)))
-
-    if args.dataset == 'Dahoas/rm-static':
-        train_dataset = RmStaticDataset(train_data, tokenizer, max_len)
-        valid_dataset = RmStaticDataset(valid_data, tokenizer, max_len)
-        eval_dataset = RmStaticDataset(eval_data, tokenizer, max_len)
-    elif args.dataset == 'Anthropic/hh-rlhf':
-        train_dataset = HhRlhfDataset(train_data, tokenizer, max_len)
-        valid_dataset = HhRlhfDataset(valid_data, tokenizer, max_len)
-        eval_dataset = HhRlhfDataset(eval_data, tokenizer, max_len)
+    train_data = data["train"].select(range(min(args.max_datasets_size, len(data["train"]))))
+    eval_data = data["test"].select(range(min(args.max_datasets_size, len(data["test"]))))
+
+    if args.dataset == "Dahoas/rm-static":
+        train_dataset = RmStaticDataset(train_data, tokenizer, args.max_len)
+        eval_dataset = RmStaticDataset(eval_data, tokenizer, args.max_len)
+    elif args.dataset == "Anthropic/hh-rlhf":
+        train_dataset = HhRlhfDataset(train_data, tokenizer, args.max_len)
+        eval_dataset = HhRlhfDataset(eval_data, tokenizer, args.max_len)
     else:
         raise ValueError(f'Unsupported dataset "{args.dataset}"')
 
     if dist.is_initialized() and dist.get_world_size() > 1:
-        train_sampler = DistributedSampler(train_dataset,
-                                           shuffle=True,
-                                           seed=42,
-                                           drop_last=True,
-                                           rank=dist.get_rank(),
-                                           num_replicas=dist.get_world_size())
-        valid_sampler = DistributedSampler(valid_dataset,
-                                           shuffle=True,
-                                           seed=42,
-                                           drop_last=True,
-                                           rank=dist.get_rank(),
-                                           num_replicas=dist.get_world_size())
-        eval_sampler = DistributedSampler(eval_dataset,
-                                          shuffle=True,
-                                          seed=42,
-                                          drop_last=True,
-                                          rank=dist.get_rank(),
-                                          num_replicas=dist.get_world_size())
+        train_sampler = DistributedSampler(
+            train_dataset,
+            shuffle=True,
+            seed=42,
+            drop_last=True,
+            rank=dist.get_rank(),
+            num_replicas=dist.get_world_size(),
+        )
+        eval_sampler = DistributedSampler(
+            eval_dataset,
+            shuffle=True,
+            seed=42,
+            drop_last=True,
+            rank=dist.get_rank(),
+            num_replicas=dist.get_world_size(),
+        )
     else:
         train_sampler = None
-        valid_sampler = None
         eval_sampler = None
 
-    train_dataloader = DataLoader(train_dataset,
-                                  shuffle=(train_sampler is None),
-                                  sampler=train_sampler,
-                                  batch_size=args.batch_size,
-                                  pin_memory=True)
-
-    valid_dataloader = DataLoader(valid_dataset,
-                                  shuffle=(valid_sampler is None),
-                                  sampler=valid_sampler,
-                                  batch_size=args.batch_size,
-                                  pin_memory=True)
-
-    eval_dataloader = DataLoader(eval_dataset,
-                                 shuffle=(eval_sampler is None),
-                                 sampler=eval_sampler,
-                                 batch_size=args.batch_size,
-                                 pin_memory=True)
-
-    (model, optim) = strategy.prepare((model, optim))
-    trainer = RewardModelTrainer(model=model,
-                                 strategy=strategy,
-                                 optim=optim,
-                                 loss_fn=loss_fn,
-                                 train_dataloader=train_dataloader,
-                                 valid_dataloader=valid_dataloader,
-                                 eval_dataloader=eval_dataloader,
-                                 max_epochs=args.max_epochs)
-
-    trainer.fit()
+    train_dataloader = DataLoader(
+        train_dataset,
+        shuffle=(train_sampler is None),
+        sampler=train_sampler,
+        batch_size=args.batch_size,
+        pin_memory=True,
+    )
+
+    eval_dataloader = DataLoader(
+        eval_dataset, shuffle=(eval_sampler is None), sampler=eval_sampler, batch_size=args.batch_size, pin_memory=True
+    )
+
+    lr_scheduler = CosineAnnealingLR(optim, train_dataloader.__len__() // 100)
+    strategy_dict = strategy.prepare(dict(model=model, optimizer=optim, lr_scheduler=lr_scheduler))
+    model = strategy_dict["model"]
+    optim = strategy_dict["optimizer"]
+    lr_scheduler = strategy_dict["lr_scheduler"]
+    trainer = RewardModelTrainer(
+        model=model,
+        strategy=strategy,
+        optim=optim,
+        lr_scheduler=lr_scheduler,
+        loss_fn=loss_fn,
+        max_epochs=args.max_epochs,
+    )
+
+    trainer.fit(
+        train_dataloader=train_dataloader,
+        eval_dataloader=eval_dataloader,
+        log_dir=args.log_dir,
+        use_wandb=args.use_wandb,
+    )
+
+    if args.lora_rank > 0 and args.merge_lora_weights:
+        from coati.models.lora import LORA_MANAGER
+
+        # NOTE: set model to eval to merge LoRA weights
+        LORA_MANAGER.merge_weights = True
+        model.eval()
     # save model checkpoint after fitting on only rank0
-    strategy.save_model(model, args.save_path, only_rank0=True)
+    state_dict = model.state_dict()
+    torch.save(state_dict, args.save_path)
     # save optimizer checkpoint on all ranks
     if args.need_optim_ckpt:
-        strategy.save_optimizer(trainer.optimizer,
-                                'rm_optim_checkpoint_%d.pt' % (torch.cuda.current_device()),
-                                only_rank0=False)
+        strategy.save_optimizer(
+            trainer.optimizer, "rm_optim_checkpoint_%d.pt" % (torch.cuda.current_device()), only_rank0=False
+        )
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    parser.add_argument('--strategy',
-                        choices=['naive', 'ddp', 'colossalai_gemini', 'colossalai_zero2'],
-                        default='colossalai_zero2')
-    parser.add_argument('--model', choices=['gpt2', 'bloom', 'opt', 'deberta', 'llama', 'roberta'], default='bloom')
-    parser.add_argument('--pretrain', type=str, default=None)
-    parser.add_argument('--model_path', type=str, default=None)
-    parser.add_argument('--need_optim_ckpt', type=bool, default=False)
-    parser.add_argument('--dataset',
-                        type=str,
-                        choices=['Anthropic/hh-rlhf', 'Dahoas/rm-static'],
-                        default='Dahoas/rm-static')
-    parser.add_argument('--subset', type=str, default=None)
-    parser.add_argument('--save_path', type=str, default='rm_ckpt')
-    parser.add_argument('--max_epochs', type=int, default=1)
-    parser.add_argument('--batch_size', type=int, default=1)
-    parser.add_argument('--max_len', type=int, default=512)
-    parser.add_argument('--lora_rank', type=int, default=0, help="low-rank adaptation matrices rank")
-    parser.add_argument('--loss_fn', type=str, default='log_sig', choices=['log_sig', 'log_exp'])
-    parser.add_argument('--test', type=bool, default=False)
+    parser.add_argument(
+        "--strategy", choices=["ddp", "colossalai_gemini", "colossalai_zero2"], default="colossalai_zero2"
+    )
+    parser.add_argument("--model", choices=["gpt2", "bloom", "opt", "llama"], default="bloom")
+    parser.add_argument("--tokenizer", type=str, default=None)
+    parser.add_argument("--pretrain", type=str, default=None)
+    parser.add_argument("--model_path", type=str, default=None)
+    parser.add_argument("--need_optim_ckpt", type=bool, default=False)
+    parser.add_argument(
+        "--dataset", type=str, choices=["Anthropic/hh-rlhf", "Dahoas/rm-static"], default="Dahoas/rm-static"
+    )
+    parser.add_argument("--subset", type=lambda x: None if x == "None" else x, default=None)
+    parser.add_argument("--max_datasets_size", type=int, default=1000000)
+    parser.add_argument("--save_path", type=str, default="rm_ckpt")
+    parser.add_argument("--max_epochs", type=int, default=1)
+    parser.add_argument("--batch_size", type=int, default=1)
+    parser.add_argument("--max_len", type=int, default=512)
+    parser.add_argument("--lora_rank", type=int, default=0, help="low-rank adaptation matrices rank")
+    parser.add_argument("--merge_lora_weights", type=bool, default=True)
+    parser.add_argument("--lr", type=float, default=9e-6)
+    parser.add_argument("--loss_fn", type=str, default="log_sig", choices=["log_sig", "log_exp"])
+    parser.add_argument("--log_dir", default="logs", type=str)
+    parser.add_argument("--use_wandb", default=False, action="store_true")
     args = parser.parse_args()
     train(args)

applications/Chat/examples/train_rm.sh

 set_n_least_used_CUDA_VISIBLE_DEVICES() {
     local n=${1:-"9999"}
     echo "GPU Memory Usage:"
-    local FIRST_N_GPU_IDS=$(nvidia-smi --query-gpu=memory.used --format=csv \
-        | tail -n +2 \
-        | nl -v 0 \
-        | tee /dev/tty \
-        | sort -g -k 2 \
-        | awk '{print $1}' \
-        | head -n $n)
+    local FIRST_N_GPU_IDS=$(nvidia-smi --query-gpu=memory.used --format=csv |
+        tail -n +2 |
+        nl -v 0 |
+        tee /dev/tty |
+        sort -g -k 2 |
+        awk '{print $1}' |
+        head -n $n)
     export CUDA_VISIBLE_DEVICES=$(echo $FIRST_N_GPU_IDS | sed 's/ /,/g')
     echo "Now CUDA_VISIBLE_DEVICES is set to:"
     echo "CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES"
@@ -16,9 +16,10 @@ set_n_least_used_CUDA_VISIBLE_DEVICES() {
 set_n_least_used_CUDA_VISIBLE_DEVICES 2
 
 torchrun --standalone --nproc_per_node=2 train_reward_model.py \
-   --pretrain  <your pretrain path> \
-   --model 'bloom' \
-   --strategy colossalai_zero2 \
-   --loss_fn 'log_sig'\
-   --save_path <your model saving path>\
-   --dataset 'Anthropic/hh-rlhf'\
+    --pretrain 'gpt2' \
+    --model 'gpt2' \
+    --strategy colossalai_zero2 \
+    --loss_fn 'log_exp' \
+    --dataset 'Anthropic/hh-rlhf' \
+    --batch_size 16 \
+    --max_epochs 10

applications/Chat/examples/train_sft.py

 import argparse
-import os
+import math
+import warnings
 
-import loralib as lora
 import torch
 import torch.distributed as dist
-from coati.dataset import DataCollatorForSupervisedDataset, SFTDataset, SupervisedDataset
-from coati.models import convert_to_lora_module
+from coati.dataset import SFTDataset, SupervisedDataset
+from coati.models.bloom import BLOOMActor
+from coati.models.chatglm import ChatGLMActor
+from coati.models.chatglm.chatglm_tokenizer import ChatGLMTokenizer
+from coati.models.gpt import GPTActor
+from coati.models.llama import LlamaActor
+from coati.models.opt import OPTActor
 from coati.trainer import SFTTrainer
-from coati.trainer.strategies import ColossalAIStrategy, DDPStrategy, NaiveStrategy
-from coati.utils import prepare_llama_tokenizer_and_embedding
+from coati.trainer.strategies import DDPStrategy, GeminiStrategy, LowLevelZeroStrategy
 from datasets import load_dataset
 from torch.optim import Adam
 from torch.utils.data import DataLoader
 from torch.utils.data.distributed import DistributedSampler
-from transformers import AutoTokenizer, BloomConfig, BloomForCausalLM, BloomTokenizerFast, LlamaConfig, LlamaForCausalLM
-from transformers.models.gpt2.configuration_gpt2 import GPT2Config
-from transformers.models.gpt2.modeling_gpt2 import GPT2LMHeadModel
+from transformers import AutoTokenizer, BloomTokenizerFast, LlamaTokenizer
 from transformers.models.gpt2.tokenization_gpt2 import GPT2Tokenizer
-from transformers.models.opt.configuration_opt import OPTConfig
-from transformers.models.opt.modeling_opt import OPTForCausalLM
+from transformers.trainer import get_scheduler
 
 from colossalai.logging import get_dist_logger
 from colossalai.nn.optimizer import HybridAdam
-from colossalai.tensor import ColoParameter
 
 
 def train(args):
     # configure strategy
-    if args.strategy == 'naive':
-        strategy = NaiveStrategy()
-    elif args.strategy == 'ddp':
+    if args.strategy == "ddp":
         strategy = DDPStrategy()
-    elif args.strategy == 'colossalai_gemini':
-        raise NotImplementedError(
-            'Gemini is not supported .from_pretrained() yet. We will update this after checkpoint io is ready.')
-        strategy = ColossalAIStrategy(stage=3, placement_policy='cuda')
-    elif args.strategy == 'colossalai_zero2':
-        strategy = ColossalAIStrategy(stage=2, placement_policy='cuda')
-    elif args.strategy == 'colossalai_zero2_cpu':
-        strategy = ColossalAIStrategy(stage=2, placement_policy='cpu')
+    elif args.strategy == "colossalai_gemini":
+        strategy = GeminiStrategy(placement_policy="auto")
+    elif args.strategy == "colossalai_zero2":
+        strategy = LowLevelZeroStrategy(stage=2, placement_policy="cuda")
+    elif args.strategy == "colossalai_zero2_cpu":
+        strategy = LowLevelZeroStrategy(stage=2, placement_policy="cpu")
     else:
         raise ValueError(f'Unsupported strategy "{args.strategy}"')
 
     # configure model
+    if args.lora_rank > 0:
+        warnings.warn("Lora is not supported yet.")
+        args.lora_rank = 0
+
     with strategy.model_init_context():
-        if args.model == 'bloom':
-            model = convert_to_lora_module(BloomForCausalLM.from_pretrained(args.pretrain),
-                                           args.lora_rank).half().cuda()
-        elif args.model == 'opt':
-            model = convert_to_lora_module(OPTForCausalLM.from_pretrained(args.pretrain), args.lora_rank).half().cuda()
-        elif args.model == 'gpt2':
-            model = convert_to_lora_module(GPT2LMHeadModel.from_pretrained(args.pretrain), args.lora_rank).half().cuda()
-        elif args.model == 'llama':
-            model = convert_to_lora_module(LlamaForCausalLM.from_pretrained(args.pretrain),
-                                           args.lora_rank).half().cuda()
+        if args.model == "bloom":
+            model = BLOOMActor(pretrained=args.pretrain, lora_rank=args.lora_rank, checkpoint=args.grad_checkpoint)
+        elif args.model == "opt":
+            model = OPTActor(pretrained=args.pretrain, lora_rank=args.lora_rank, checkpoint=args.grad_checkpoint)
+        elif args.model == "gpt2":
+            model = GPTActor(pretrained=args.pretrain, lora_rank=args.lora_rank, checkpoint=args.grad_checkpoint)
+        elif args.model == "llama":
+            model = LlamaActor(pretrained=args.pretrain, lora_rank=args.lora_rank, checkpoint=args.grad_checkpoint)
+        elif args.model == "chatglm":
+            model = ChatGLMActor(pretrained=args.pretrain)
         else:
             raise ValueError(f'Unsupported model "{args.model}"')
-    if args.grad_checkpoint:
-        model.gradient_checkpointing_enable()
+
+        model.to(torch.bfloat16).to(torch.cuda.current_device())
 
     # configure tokenizer
-    if args.model == 'gpt2':
-        tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
+    if args.model == "gpt2":
+        tokenizer = GPT2Tokenizer.from_pretrained("gpt2" if args.tokenizer is None else args.tokenizer)
+        tokenizer.pad_token = tokenizer.eos_token
+    elif args.model == "bloom":
+        tokenizer = BloomTokenizerFast.from_pretrained(
+            "bigscience/bloom-560m" if args.tokenizer is None else args.tokenizer
+        )
         tokenizer.pad_token = tokenizer.eos_token
-    elif args.model == 'bloom':
-        tokenizer = BloomTokenizerFast.from_pretrained(args.pretrain)
+    elif args.model == "opt":
+        tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m" if args.tokenizer is None else args.tokenizer)
         tokenizer.pad_token = tokenizer.eos_token
-    elif args.model == 'opt':
-        tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m")
-    elif args.model == 'llama':
-        tokenizer = AutoTokenizer.from_pretrained(
-            args.pretrain,
-            padding_side="right",
-            use_fast=False,
+    elif args.model == "llama":
+        tokenizer = LlamaTokenizer.from_pretrained(
+            "hf-internal-testing/llama-tokenizer" if args.tokenizer is None else args.tokenizer
+        )
+        tokenizer.eos_token = "<\s>"
+        tokenizer.pad_token = tokenizer.unk_token
+    elif args.model == "chatglm":
+        tokenizer = ChatGLMTokenizer.from_pretrained(
+            "THUDM/chatglm-6b" if args.tokenizer is None else args.tokenizer, trust_remote_code=True
         )
-        tokenizer.eos_token = '<\s>'
     else:
         raise ValueError(f'Unsupported model "{args.model}"')
-    tokenizer.pad_token = tokenizer.eos_token
-    max_len = args.max_len
-    if args.model == 'llama':
-        tokenizer = prepare_llama_tokenizer_and_embedding(tokenizer, model)
-
-        if args.strategy == 'colossalai_gemini':
-            # this is a hack to deal with the resized embedding
-            # to make sure all parameters are ColoParameter for Colossal-AI Gemini Compatiblity
-            for name, param in model.named_parameters():
-                if not isinstance(param, ColoParameter):
-                    sub_module_name = '.'.join(name.split('.')[:-1])
-                    weight_name = name.split('.')[-1]
-                    sub_module = model.get_submodule(sub_module_name)
-                    setattr(sub_module, weight_name, ColoParameter(param))
-    else:
-        tokenizer.pad_token = tokenizer.eos_token
 
     # configure optimizer
-    if args.strategy.startswith('colossalai'):
+    if args.strategy.startswith("colossalai"):
         optim = HybridAdam(model.parameters(), lr=args.lr, clipping_norm=1.0)
     else:
         optim = Adam(model.parameters(), lr=args.lr)
 
-    logger = get_dist_logger()
-
     # configure dataset
-    if args.dataset == 'yizhongw/self_instruct':
-        train_data = load_dataset(args.dataset, 'super_natural_instructions', split='train')
-        eval_data = load_dataset(args.dataset, 'super_natural_instructions', split='test')
+    if args.dataset == "yizhongw/self_instruct":
+        train_data = load_dataset(args.dataset, "super_natural_instructions", split="train")
+        eval_data = load_dataset(args.dataset, "super_natural_instructions", split="test")
+
+        if args.max_datasets_size is not None:
+            train_data = train_data.select(range(min(args.max_datasets_size, len(train_data))))
+            eval_data = eval_data.select(range(min(args.max_datasets_size, len(eval_data))))
 
-        train_dataset = SFTDataset(train_data, tokenizer, max_len)
-        eval_dataset = SFTDataset(eval_data, tokenizer, max_len)
+        train_dataset = SFTDataset(train_data, tokenizer, args.max_len)
+        eval_dataset = SFTDataset(eval_data, tokenizer, args.max_len)
 
     else:
-        train_dataset = SupervisedDataset(tokenizer=tokenizer,
-                                          data_path=args.dataset,
-                                          max_datasets_size=args.max_datasets_size,
-                                          max_length=max_len)
+        train_dataset = SupervisedDataset(
+            tokenizer=tokenizer,
+            data_path=args.dataset,
+            max_datasets_size=args.max_datasets_size,
+            max_length=args.max_len,
+        )
         eval_dataset = None
-    data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)
 
     if dist.is_initialized() and dist.get_world_size() > 1:
-        train_sampler = DistributedSampler(train_dataset,
-                                           shuffle=True,
-                                           seed=42,
-                                           drop_last=True,
-                                           rank=dist.get_rank(),
-                                           num_replicas=dist.get_world_size())
+        train_sampler = DistributedSampler(
+            train_dataset,
+            shuffle=True,
+            seed=42,
+            drop_last=True,
+            rank=dist.get_rank(),
+            num_replicas=dist.get_world_size(),
+        )
         if eval_dataset is not None:
-            eval_sampler = DistributedSampler(eval_dataset,
-                                              shuffle=False,
-                                              seed=42,
-                                              drop_last=False,
-                                              rank=dist.get_rank(),
-                                              num_replicas=dist.get_world_size())
+            eval_sampler = DistributedSampler(
+                eval_dataset,
+                shuffle=False,
+                seed=42,
+                drop_last=False,
+                rank=dist.get_rank(),
+                num_replicas=dist.get_world_size(),
+            )
     else:
         train_sampler = None
         eval_sampler = None
 
-    train_dataloader = DataLoader(train_dataset,
-                                  shuffle=(train_sampler is None),
-                                  sampler=train_sampler,
-                                  batch_size=args.batch_size,
-                                  collate_fn=data_collator,
-                                  pin_memory=True)
+    train_dataloader = DataLoader(
+        train_dataset,
+        shuffle=(train_sampler is None),
+        sampler=train_sampler,
+        batch_size=args.batch_size,
+        pin_memory=True,
+    )
     if eval_dataset is not None:
-        eval_dataloader = DataLoader(eval_dataset,
-                                     shuffle=(eval_sampler is None),
-                                     sampler=eval_sampler,
-                                     batch_size=args.batch_size,
-                                     collate_fn=data_collator,
-                                     pin_memory=True)
+        eval_dataloader = DataLoader(
+            eval_dataset,
+            shuffle=(eval_sampler is None),
+            sampler=eval_sampler,
+            batch_size=args.batch_size,
+            pin_memory=True,
+        )
     else:
         eval_dataloader = None
 
-    (model, optim) = strategy.prepare((model, optim))
-    trainer = SFTTrainer(model=model,
-                         strategy=strategy,
-                         optim=optim,
-                         train_dataloader=train_dataloader,
-                         eval_dataloader=eval_dataloader,
-                         max_epochs=args.max_epochs,
-                         accumulation_steps=args.accumulation_steps)
-
-    trainer.fit(logger=logger, use_wandb=args.use_wandb)
+    num_update_steps_per_epoch = len(train_dataloader) // args.accumulation_steps
+    max_steps = math.ceil(args.max_epochs * num_update_steps_per_epoch)
+    lr_scheduler = get_scheduler(
+        "cosine", optim, num_warmup_steps=math.ceil(max_steps * 0.03), num_training_steps=max_steps
+    )
+    strategy_dict = strategy.prepare(dict(model=model, optimizer=optim, lr_scheduler=lr_scheduler))
+    model = strategy_dict["model"]
+    optim = strategy_dict["optimizer"]
+    lr_scheduler = strategy_dict["lr_scheduler"]
+    trainer = SFTTrainer(
+        model=model,
+        strategy=strategy,
+        optim=optim,
+        lr_scheduler=lr_scheduler,
+        max_epochs=args.max_epochs,
+        accumulation_steps=args.accumulation_steps,
+    )
 
+    logger = get_dist_logger()
+    trainer.fit(
+        train_dataloader=train_dataloader,
+        eval_dataloader=eval_dataloader,
+        logger=logger,
+        log_dir=args.log_dir,
+        use_wandb=args.use_wandb,
+    )
+
+    if args.lora_rank > 0 and args.merge_lora_weights:
+        from coati.models.lora import LORA_MANAGER
+
+        # NOTE: set model to eval to merge LoRA weights
+        LORA_MANAGER.merge_weights = True
+        model.eval()
     # save model checkpoint after fitting on only rank0
-    strategy.save_pretrained(model, path=args.save_path, only_rank0=True, tokenizer=tokenizer)
+    strategy.save_pretrained(model, path=args.save_path, tokenizer=tokenizer)
     # save optimizer checkpoint on all ranks
     if args.need_optim_ckpt:
-        strategy.save_optimizer(trainer.optimizer,
-                                'rm_optim_checkpoint_%d.pt' % (torch.cuda.current_device()),
-                                only_rank0=False)
+        strategy.save_optimizer(
+            trainer.optimizer, "rm_optim_checkpoint_%d.pt" % (torch.cuda.current_device()), only_rank0=False
+        )
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    parser.add_argument('--strategy',
-                        choices=['naive', 'ddp', 'colossalai_gemini', 'colossalai_zero2', 'colossalai_zero2_cpu'],
-                        default='colossalai_zero2')
-    parser.add_argument('--model', choices=['gpt2', 'bloom', 'opt', 'llama'], default='bloom')
-    parser.add_argument('--pretrain', type=str, default=None)
-    parser.add_argument('--dataset', type=str, default=None)
-    parser.add_argument('--max_datasets_size', type=int, default=None)
-    parser.add_argument('--save_path', type=str, default='output')
-    parser.add_argument('--need_optim_ckpt', type=bool, default=False)
-    parser.add_argument('--max_epochs', type=int, default=3)
-    parser.add_argument('--batch_size', type=int, default=4)
-    parser.add_argument('--max_len', type=int, default=512)
-    parser.add_argument('--lora_rank', type=int, default=0, help="low-rank adaptation matrices rank")
-    parser.add_argument('--log_interval', type=int, default=100, help="how many steps to log")
-    parser.add_argument('--lr', type=float, default=5e-6)
-    parser.add_argument('--accumulation_steps', type=int, default=8)
-    parser.add_argument('--use_wandb', default=False, action='store_true')
-    parser.add_argument('--grad_checkpoint', default=False, action='store_true')
+    parser.add_argument(
+        "--strategy",
+        choices=["ddp", "colossalai_gemini", "colossalai_zero2", "colossalai_zero2_cpu"],
+        default="colossalai_zero2",
+    )
+    parser.add_argument("--model", choices=["gpt2", "bloom", "opt", "llama", "chatglm"], default="bloom")
+    parser.add_argument("--tokenizer", type=str, default=None)
+    parser.add_argument("--pretrain", type=str, default=None)
+    parser.add_argument("--dataset", type=str, default=None)
+    parser.add_argument("--max_datasets_size", type=int, default=None)
+    parser.add_argument("--save_path", type=str, default="output")
+    parser.add_argument("--need_optim_ckpt", type=bool, default=False)
+    parser.add_argument("--max_epochs", type=int, default=3)
+    parser.add_argument("--batch_size", type=int, default=4)
+    parser.add_argument("--max_len", type=int, default=512)
+    parser.add_argument("--lora_rank", type=int, default=0, help="low-rank adaptation matrices rank")
+    parser.add_argument("--merge_lora_weights", type=bool, default=True)
+    parser.add_argument("--lr", type=float, default=5e-6)
+    parser.add_argument("--accumulation_steps", type=int, default=8)
+    parser.add_argument("--log_dir", default="logs", type=str)
+    parser.add_argument("--use_wandb", default=False, action="store_true")
+    parser.add_argument("--grad_checkpoint", default=False, action="store_true")
     args = parser.parse_args()
     train(args)

applications/Chat/examples/train_sft.sh

+set_n_least_used_CUDA_VISIBLE_DEVICES() {
+    local n=${1:-"9999"}
+    echo "GPU Memory Usage:"
+    local FIRST_N_GPU_IDS=$(nvidia-smi --query-gpu=memory.used --format=csv |
+        tail -n +2 |
+        nl -v 0 |
+        tee /dev/tty |
+        sort -g -k 2 |
+        awk '{print $1}' |
+        head -n $n)
+    export CUDA_VISIBLE_DEVICES=$(echo $FIRST_N_GPU_IDS | sed 's/ /,/g')
+    echo "Now CUDA_VISIBLE_DEVICES is set to:"
+    echo "CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES"
+}
+
+set_n_least_used_CUDA_VISIBLE_DEVICES 4
+
 torchrun --standalone --nproc_per_node=4 train_sft.py \
     --pretrain "/path/to/LLaMa-7B/" \
     --model 'llama' \
     --strategy colossalai_zero2 \
-    --log_interval 10 \
-    --save_path  /path/to/Coati-7B \
+    --save_path /path/to/Coati-7B \
     --dataset /path/to/data.json \
     --batch_size 4 \
     --accumulation_steps 8 \
     --lr 2e-5 \
     --max_datasets_size 512 \
-    --max_epochs 1 \
+    --max_epochs 1

applications/Chat/inference/README.md

@@ -20,21 +20,21 @@ Tha data is from [LLaMA Int8 4bit ChatBot Guide v2](https://rentry.org/llama-tar
 
 ### 8-bit
 
-| Model | Min GPU RAM | Recommended GPU RAM | Min RAM/Swap | Card examples |
-| :---: | :---: | :---: | :---: | :---: |
-| LLaMA-7B | 9.2GB | 10GB | 24GB | 3060 12GB, RTX 3080 10GB, RTX 3090 |
-| LLaMA-13B | 16.3GB | 20GB | 32GB | RTX 3090 Ti, RTX 4090 |
-| LLaMA-30B | 36GB | 40GB | 64GB | A6000 48GB, A100 40GB |
-| LLaMA-65B | 74GB | 80GB | 128GB | A100 80GB |
+|   Model   | Min GPU RAM | Recommended GPU RAM | Min RAM/Swap |           Card examples            |
+| :-------: | :---------: | :-----------------: | :----------: | :--------------------------------: |
+| LLaMA-7B  |    9.2GB    |        10GB         |     24GB     | 3060 12GB, RTX 3080 10GB, RTX 3090 |
+| LLaMA-13B |   16.3GB    |        20GB         |     32GB     |       RTX 3090 Ti, RTX 4090        |
+| LLaMA-30B |    36GB     |        40GB         |     64GB     |       A6000 48GB, A100 40GB        |
+| LLaMA-65B |    74GB     |        80GB         |    128GB     |             A100 80GB              |
 
 ### 4-bit
 
-| Model | Min GPU RAM | Recommended GPU RAM | Min RAM/Swap | Card examples |
-| :---: | :---: | :---: | :---: | :---: |
-| LLaMA-7B | 3.5GB | 6GB | 16GB | RTX 1660, 2060, AMD 5700xt, RTX 3050, 3060 |
-| LLaMA-13B | 6.5GB | 10GB | 32GB | AMD 6900xt, RTX 2060 12GB, 3060 12GB, 3080, A2000 |
-| LLaMA-30B | 15.8GB | 20GB | 64GB | RTX 3080 20GB, A4500, A5000, 3090, 4090, 6000, Tesla V100 |
-| LLaMA-65B | 31.2GB | 40GB | 128GB | A100 40GB, 2x3090, 2x4090, A40, RTX A6000, 8000, Titan Ada |
+|   Model   | Min GPU RAM | Recommended GPU RAM | Min RAM/Swap |                       Card examples                        |
+| :-------: | :---------: | :-----------------: | :----------: | :--------------------------------------------------------: |
+| LLaMA-7B  |    3.5GB    |         6GB         |     16GB     |         RTX 1660, 2060, AMD 5700xt, RTX 3050, 3060         |
+| LLaMA-13B |    6.5GB    |        10GB         |     32GB     |     AMD 6900xt, RTX 2060 12GB, 3060 12GB, 3080, A2000      |
+| LLaMA-30B |   15.8GB    |        20GB         |     64GB     | RTX 3080 20GB, A4500, A5000, 3090, 4090, 6000, Tesla V100  |
+| LLaMA-65B |   31.2GB    |        40GB         |    128GB     | A100 40GB, 2x3090, 2x4090, A40, RTX A6000, 8000, Titan Ada |
 
 ## General setup
 
@@ -75,7 +75,7 @@ E.g. you can set `export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH`.
 
 Please ensure you have downloaded HF-format model weights of LLaMA models first.
 
-Then you can follow [GPTQ-for-LLaMa](https://github.com/qwopqwop200/GPTQ-for-LLaMa). This lib provides efficient CUDA kernels and weight convertion script.
+Then you can follow [GPTQ-for-LLaMa](https://github.com/qwopqwop200/GPTQ-for-LLaMa). This lib provides efficient CUDA kernels and weight conversion script.
 
 After installing this lib, we may convert the original HF-format LLaMA model weights to 4-bit version.
 

applications/Chat/inference/benchmark.py

@@ -4,8 +4,8 @@ import argparse
 from time import time
 
 import torch
-from llama_gptq import load_quant
-from transformers import AutoTokenizer, GenerationConfig, LlamaForCausalLM
+from coati.quant import llama_load_quant, low_resource_init
+from transformers import AutoTokenizer, GenerationConfig, LlamaConfig, LlamaForCausalLM
 
 
 def generate_prompt(instruction, input=None):
@@ -84,49 +84,58 @@ inst = [instructions[0]] * 4
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument(
-        'pretrained',
-        help='Path to pretrained model. Can be a local path or a model name from the HuggingFace model hub.')
-    parser.add_argument('--quant',
-                        choices=['8bit', '4bit'],
-                        default=None,
-                        help='Quantization mode. Default: None (no quantization, fp16).')
+        "pretrained",
+        help="Path to pretrained model. Can be a local path or a model name from the HuggingFace model hub.",
+    )
+    parser.add_argument(
+        "--quant",
+        choices=["8bit", "4bit"],
+        default=None,
+        help="Quantization mode. Default: None (no quantization, fp16).",
+    )
     parser.add_argument(
-        '--gptq_checkpoint',
+        "--gptq_checkpoint",
         default=None,
-        help='Path to GPTQ checkpoint. This is only useful when quantization mode is 4bit. Default: None.')
-    parser.add_argument('--gptq_group_size',
-                        type=int,
-                        default=128,
-                        help='Group size for GPTQ. This is only useful when quantization mode is 4bit. Default: 128.')
+        help="Path to GPTQ checkpoint. This is only useful when quantization mode is 4bit. Default: None.",
+    )
+    parser.add_argument(
+        "--gptq_group_size",
+        type=int,
+        default=128,
+        help="Group size for GPTQ. This is only useful when quantization mode is 4bit. Default: 128.",
+    )
     args = parser.parse_args()
 
-    if args.quant == '4bit':
-        assert args.gptq_checkpoint is not None, 'Please specify a GPTQ checkpoint.'
+    if args.quant == "4bit":
+        assert args.gptq_checkpoint is not None, "Please specify a GPTQ checkpoint."
 
     tokenizer = AutoTokenizer.from_pretrained(args.pretrained)
 
-    if args.quant == '4bit':
-        model = load_quant(args.pretrained, args.gptq_checkpoint, 4, args.gptq_group_size)
+    if args.quant == "4bit":
+        with low_resource_init():
+            config = LlamaConfig.from_pretrained(args.pretrained)
+            model = LlamaForCausalLM(config)
+        model = llama_load_quant(model, args.gptq_checkpoint, 4, args.gptq_group_size)
         model.cuda()
     else:
         model = LlamaForCausalLM.from_pretrained(
             args.pretrained,
-            load_in_8bit=(args.quant == '8bit'),
+            load_in_8bit=(args.quant == "8bit"),
             torch_dtype=torch.float16,
             device_map="auto",
         )
-        if args.quant != '8bit':
-            model.half()    # seems to fix bugs for some users.
+        if args.quant != "8bit":
+            model.half()  # seems to fix bugs for some users.
         model.eval()
 
     total_tokens = 0
     start = time()
     for instruction in instructions:
         print(f"Instruction: {instruction}")
-        resp, tokens = evaluate(model, tokenizer, instruction, temparature=0.2, num_beams=1)
+        resp, tokens = evaluate(model, tokenizer, instruction, temperature=0.2, num_beams=1)
         total_tokens += tokens
         print(f"Response: {resp}")
-        print('\n----------------------------\n')
+        print("\n----------------------------\n")
     duration = time() - start
-    print(f'Total time: {duration:.3f} s, {total_tokens/duration:.3f} tokens/s')
-    print(f'Peak CUDA mem: {torch.cuda.max_memory_allocated()/1024**3:.3f} GB')
+    print(f"Total time: {duration:.3f} s, {total_tokens/duration:.3f} tokens/s")
+    print(f"Peak CUDA mem: {torch.cuda.max_memory_allocated()/1024**3:.3f} GB")

applications/Chat/inference/locustfile.py

-from json import JSONDecodeError
-
 from locust import HttpUser, task
 
-samples = [[
-    dict(
-        instruction='Who is the best player in the history of NBA?',
-        response=
-        'The best player in the history of the NBA is widely considered to be Michael Jordan. He is one of the most successful players in the league, having won 6 NBA championships with the Chicago Bulls and 5 more with the Washington Wizards. He is a 5-time MVP, 1'
-    ),
-    dict(instruction='continue this talk', response=''),
-], [
-    dict(instruction='Who is the best player in the history of NBA?', response=''),
-]]
+samples = [
+    [
+        dict(
+            instruction="Who is the best player in the history of NBA?",
+            response="The best player in the history of the NBA is widely considered to be Michael Jordan. He is one of the most successful players in the league, having won 6 NBA championships with the Chicago Bulls and 5 more with the Washington Wizards. He is a 5-time MVP, 1",
+        ),
+        dict(instruction="continue this talk", response=""),
+    ],
+    [
+        dict(instruction="Who is the best player in the history of NBA?", response=""),
+    ],
+]
 
 
 class GenerationUser(HttpUser):
-
     @task
     def generate(self):
         for sample in samples:
-            data = {'max_new_tokens': 64, 'history': sample}
-            with self.client.post('/generate', json=data, catch_response=True) as response:
+            data = {"max_new_tokens": 64, "history": sample}
+            with self.client.post("/generate", json=data, catch_response=True) as response:
                 if response.status_code in (200, 406):
                     response.success()
                 else:
-                    response.failure('Response wrong')
+                    response.failure("Response wrong")

applications/Chat/inference/requirements.txt

@@ -10,4 +10,4 @@ uvicorn
 git+https://github.com/huggingface/transformers
 accelerate
 bitsandbytes
-jieba
\ No newline at end of file
+jieba