Repository: hpcaitech/ColossalAI
File: applications/Chat/benchmarks/benchmark_opt_lora_dummy.py (Python)

Code before:


def get_gpt_config(model_name: str) -> OPTConfig:
    model_map = {
        "125m": OPTConfig.from_pretrained("facebook/opt-125m"),
        "350m": OPTConfig(hidden_size=1024, ffn_dim=4096, num_hidden_layers=24, num_attention_heads=16),
        "700m": OPTConfig(hidden_size=1280, ffn_dim=5120, num_hidden_layers=36, num_attention_heads=20),
        "1.3b": OPTConfig.from_pretrained("facebook/opt-1.3b"),
        "2.7b": OPTConfig.from_pretrained("facebook/opt-2.7b"),
        "3.5b": OPTConfig(hidden_size=3072, ffn_dim=12288, num_hidden_layers=32, num_attention_heads=32),
        "5.5b": OPTConfig(hidden_size=3840, ffn_dim=15360, num_hidden_layers=32, num_attention_heads=32),
        "6.7b": OPTConfig.from_pretrained("facebook/opt-6.7b"),
        "10b": OPTConfig(hidden_size=5120, ffn_dim=20480, num_hidden_layers=32, num_attention_heads=32),
        "13b": OPTConfig.from_pretrained("facebook/opt-13b"),
    }
    try:
        return model_map[model_name]
    except KeyError:
        raise ValueError(f'Unknown model "{model_name}"')


def main(args):
    if args.strategy == "ddp":
        strategy = DDPStrategy()
    elif args.strategy == "colossalai_gemini":
        strategy = GeminiStrategy(placement_policy="auto", initial_scale=2**5)
    elif args.strategy == "colossalai_gemini_cpu":
        strategy = GeminiStrategy(placement_policy="static", initial_scale=2**5)
    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")
    elif args.strategy == "colossalai_zero1":
        strategy = LowLevelZeroStrategy(stage=1, placement_policy="cuda")
    elif args.strategy == "colossalai_zero1_cpu":
        strategy = LowLevelZeroStrategy(stage=1, placement_policy="cpu")
    else:
        raise ValueError(f'Unsupported strategy "{args.strategy}"')

    torch.cuda.set_per_process_memory_fraction(args.cuda_mem_frac)

    model_config = get_gpt_config(args.model)
    critic_config = get_gpt_config(args.critic_model)
    with strategy.model_init_context():
        actor = OPTActor(config=model_config, lora_rank=args.lora_rank).cuda()
        critic = OPTCritic(config=critic_config, lora_rank=args.lora_rank).cuda()

        initial_model = deepcopy(actor).cuda().half()
        reward_model = RewardModel(deepcopy(critic.model), deepcopy(critic.value_head)).cuda().half()

    if args.use_kernels:

Review comment (suggestion):
This is not equivalence. Bounding `"static"` to `"colossalai_gemini"` and `"auto"` to `"colossalai_gemini_auto"` is better.

Code after:


def get_gpt_config(model_name: str) -> OPTConfig:
    model_map = {
        "125m": OPTConfig.from_pretrained("facebook/opt-125m"),
        "350m": OPTConfig(hidden_size=1024, ffn_dim=4096, num_hidden_layers=24, num_attention_heads=16),
        "700m": OPTConfig(hidden_size=1280, ffn_dim=5120, num_hidden_layers=36, num_attention_heads=20),
        "1.3b": OPTConfig.from_pretrained("facebook/opt-1.3b"),
        "2.7b": OPTConfig.from_pretrained("facebook/opt-2.7b"),
        "3.5b": OPTConfig(hidden_size=3072, ffn_dim=12288, num_hidden_layers=32, num_attention_heads=32),
        "5.5b": OPTConfig(hidden_size=3840, ffn_dim=15360, num_hidden_layers=32, num_attention_heads=32),
        "6.7b": OPTConfig.from_pretrained("facebook/opt-6.7b"),
        "10b": OPTConfig(hidden_size=5120, ffn_dim=20480, num_hidden_layers=32, num_attention_heads=32),
        "13b": OPTConfig.from_pretrained("facebook/opt-13b"),
    }
    try:
        return model_map[model_name]
    except KeyError:
        raise ValueError(f'Unknown model "{model_name}"')


def main(args):
    if args.strategy == "ddp":
        strategy = DDPStrategy()
    elif args.strategy == "colossalai_gemini":
        strategy = GeminiStrategy(placement_policy="static",initial_scale=2**5)
    elif args.strategy == "colossalai_gemini_cpu":
        strategy = GeminiStrategy(placement_policy="static", offload_optim_frac=1.0, offload_param_frac=1.0, initial_scale=2**5)
    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")
    elif args.strategy == "colossalai_zero1":
        strategy = LowLevelZeroStrategy(stage=1, placement_policy="cuda")
    elif args.strategy == "colossalai_zero1_cpu":
        strategy = LowLevelZeroStrategy(stage=1, placement_policy="cpu")
    else:
        raise ValueError(f'Unsupported strategy "{args.strategy}"')

    torch.cuda.set_per_process_memory_fraction(args.cuda_mem_frac)

    model_config = get_gpt_config(args.model)
    critic_config = get_gpt_config(args.critic_model)
    with strategy.model_init_context():
        actor = OPTActor(config=model_config, lora_rank=args.lora_rank).cuda()
        critic = OPTCritic(config=critic_config, lora_rank=args.lora_rank).cuda()

        initial_model = deepcopy(actor).cuda().half()
        reward_model = RewardModel(deepcopy(critic.model), deepcopy(critic.value_head)).cuda().half()

    if args.use_kernels: