(args, pipe, device, train_data_loader, test_data_loader)
| 23 | |
| 24 | |
| 25 | def train_loop(args, pipe, device, train_data_loader, test_data_loader): |
| 26 | |
| 27 | print('training starts......') |
| 28 | |
| 29 | pipe.model.train() # Flag .training to True to enable Dropout |
| 30 | |
| 31 | use_dp = (args.world_size != args.pipeline_group_size) |
| 32 | if use_dp: |
| 33 | # dp_comm = get_data_parallel_comm() |
| 34 | dp_rank = get_data_parallel_rank() |
| 35 | dp_size = get_data_parallel_world_size() |
| 36 | else: |
| 37 | dp_rank = 0 |
| 38 | dp_size = 1 |
| 39 | pp_comm = get_pipeline_parallel_comm() |
| 40 | |
| 41 | stop_flag = torch.zeros(1, dtype=torch.int64).to(device) |
| 42 | |
| 43 | input_ids = torch.zeros( |
| 44 | [args.batch_size, args.seq_length], |
| 45 | dtype=torch.int64 |
| 46 | ).to(device) |
| 47 | |
| 48 | prefix_masks = torch.zeros( |
| 49 | [args.batch_size, args.seq_length], |
| 50 | dtype=torch.uint8 |
| 51 | ).to(device) |
| 52 | |
| 53 | do_sync_before_save = (args.dp_mode in ['local'] and use_dp) |
| 54 | |
| 55 | if get_pipeline_parallel_rank() == 0 and dp_rank == 0: |
| 56 | |
| 57 | for i, data in enumerate(train_data_loader): |
| 58 | if i < pipe.global_step: |
| 59 | continue |
| 60 | |
| 61 | if use_dp: |
| 62 | get_data_parallel_comm().broadcast(stop_flag, 0) |
| 63 | pp_comm.broadcast(stop_flag, 0) |
| 64 | |
| 65 | if stop_flag.item() == 1: |
| 66 | break |
| 67 | |
| 68 | input_ids_global = data['input_ids'].to(torch.int64).to(device) |
| 69 | prefix_masks_global = data['prefix_masks'].to(torch.uint8).to(device) |
| 70 | |
| 71 | input_ids_list = input_ids_global.chunk(dp_size) |
| 72 | prefix_masks_list = prefix_masks_global.chunk(dp_size) |
| 73 | |
| 74 | if use_dp: |
| 75 | for j in range(1, dp_size): |
| 76 | get_data_parallel_comm().send( |
| 77 | input_ids_list[j], j, |
| 78 | ) |
| 79 | get_data_parallel_comm().send( |
| 80 | prefix_masks_list[j], j, |
| 81 | ) |
| 82 |
no test coverage detected