Function main

data/generation/generate.py:149–237

(rank, args)

    return gen_dataset, data_collator

def main(rank, args):
    dist.init_process_group("nccl")
    torch.manual_seed(args.seed)
    world_size = torch.cuda.device_count()

    base_model = args.base_model
    batch_size = args.batch_size
    return_seq_num = 1
    
    n_gpus = torch.cuda.device_count()

    model = AutoModelForCausalLM.from_pretrained(
            base_model,
            torch_dtype=torch.bfloat16
        )
    
    tokenizer = AutoTokenizer.from_pretrained(base_model, use_fast=False)
    tokenizer.truncation_side = 'left'
    if tokenizer.pad_token is None:
        smart_tokenizer_and_embedding_resize(
            special_tokens_dict=dict(pad_token=DEFAULT_PAD_TOKEN),
            tokenizer=tokenizer,
            model=model,
        )

    torch.cuda.set_device(rank)
    model.to(torch.cuda.current_device())
    model = DDP(model, device_ids=[torch.cuda.current_device()])
    model.eval()

    # Get the generation dataset
    gen_dataset, data_collator = make_supervised_data_module(tokenizer, args.dataset_name, args.max_sample)

    sampler = torch.utils.data.distributed.DistributedSampler(gen_dataset, num_replicas=world_size, rank=rank, shuffle=False)
    dataloader = DataLoader(
        gen_dataset, 
        shuffle=False, 
        collate_fn=data_collator, 
        batch_size=batch_size,
        sampler=sampler,
        drop_last=True
    )

    generation_config = GenerationConfig(
        temperature=args.temperature,
        do_sample=True,
        num_beams=return_seq_num,
        max_new_tokens=args.max_new_tokens,
        num_return_sequences=return_seq_num,
        top_p=1.0
    )

    all_outputs = []
    total_nums = int(len(gen_dataset) / (world_size * batch_size))
    for step, batch in tqdm(enumerate(dataloader), total=total_nums):
        input_ids = batch['input_ids'].to(model.device)
        attention_mask = batch['attention_mask'].to(model.device)
        with torch.no_grad():