(rank, args)
| 192 | |
| 193 | |
| 194 | def main(rank, args): |
| 195 | dist.init_process_group("nccl") |
| 196 | torch.manual_seed(args.seed) |
| 197 | world_size = torch.cuda.device_count() |
| 198 | base_model = args.base_model |
| 199 | data_path = args.data_path |
| 200 | batch_size = args.batch_size |
| 201 | |
| 202 | n_gpus = torch.cuda.device_count() |
| 203 | max_memory = f'80000MB' |
| 204 | max_memory = {i: max_memory for i in range(n_gpus)} |
| 205 | device_map = "auto" |
| 206 | |
| 207 | # if we are in a distributed setting, we need to set the device map and max memory per device |
| 208 | if os.environ.get('LOCAL_RANK') is not None: |
| 209 | local_rank = int(os.environ.get('LOCAL_RANK', '0')) |
| 210 | device_map = {'': local_rank} |
| 211 | max_memory = {'': max_memory[local_rank]} |
| 212 | |
| 213 | model = LlamaForCausalLM.from_pretrained( |
| 214 | base_model, |
| 215 | torch_dtype=torch.bfloat16, |
| 216 | device_map=device_map, |
| 217 | ) |
| 218 | |
| 219 | if args.quant_type is not None: |
| 220 | q_config = { |
| 221 | "zero_point": True, # by default True |
| 222 | "q_group_size": args.group_size, # whether to use group quantization |
| 223 | } |
| 224 | pseudo_quantize_model_weight( |
| 225 | model, w_bit=args.bits, q_config=q_config, quant_type=args.quant_type |
| 226 | ) |
| 227 | |
| 228 | # model.half() |
| 229 | tokenizer = transformers.AutoTokenizer.from_pretrained(base_model, use_fast=False) |
| 230 | |
| 231 | torch.cuda.set_device(rank) |
| 232 | model.to(torch.cuda.current_device()) |
| 233 | model = DDP(model, device_ids=[torch.cuda.current_device()]) |
| 234 | model.eval() |
| 235 | |
| 236 | eval_dataset, data_collator = make_supervised_data_module(tokenizer, data_path) |
| 237 | # dataset_for_eval = load_dataset(data_path)['train'] |
| 238 | return_seq_num = 1 |
| 239 | for tempera in [0.2]: |
| 240 | sampler = torch.utils.data.distributed.DistributedSampler(eval_dataset, num_replicas=world_size, rank=rank, shuffle=False) |
| 241 | dataloader = DataLoader( |
| 242 | eval_dataset, |
| 243 | shuffle=False, |
| 244 | collate_fn=data_collator, |
| 245 | batch_size=batch_size, |
| 246 | sampler=sampler, |
| 247 | drop_last=True |
| 248 | ) |
| 249 | generation_config = GenerationConfig( |
| 250 | # temperature=0.8 if args.diverse_beam > 1 else 1.0, |
| 251 | temperature=tempera, |
no test coverage detected