Function main

evaluation/generation.py:283–312

(args, node_rank: int, local_rank: int, master_port: int, num_devices: int)

def main(args, node_rank: int, local_rank: int, master_port: int, num_devices: int):
    world_size = args.gen_node_world_size * num_devices
    args.rank = num_devices * node_rank + local_rank
    args.world_size = world_size
    logger.info(f"Generating on rank {args.rank} of {args.world_size}")
    
    try:
        if args.model_name in ["codegeex2-6b"]:
            tokenizer = AutoTokenizer.from_pretrained(args.model_path, trust_remote_code=True)
        else:
            tokenizer = AutoTokenizer.from_pretrained(args.model_path, clean_up_tokenization_spaces=False, trust_remote_code=True)
        if args.model_name in ["codegeex2-6b"]:
            model = AutoModel.from_pretrained(args.model_path, trust_remote_code=True).to("cuda:{}".format(local_rank % torch.cuda.device_count()))
        elif args.model_name in ["starcoder", "replit-code-v1-3b", "codegen25-7b-multi", "codegen25-7b-mono", "codegen-16B-multi"]:
            model = AutoModelForCausalLM.from_pretrained(args.model_path, trust_remote_code=True).to("cuda:{}".format(local_rank % torch.cuda.device_count()))
        else:
            try:
                model = AutoModel.from_pretrained(args.model_path, trust_remote_code=True).to("cuda:{}".format(local_rank % torch.cuda.device_count()))
            except:
                logger.error(f"Model {args.model_name} not supported.")
                raise NotImplementedError
    except Exception as e:
        logger.error(e)
    
    model = model.eval()
    # Generate samples.
    run_generation_distributed(args, model, tokenizer)

    logger.info(f"rank={args.rank} worker finished, waiting ...")
    exit(0)


def server(args):