已合并
extra token优化, mtp layer重计算和tp支持 #2242
AtomGit-Bot创建于 2025年2月18日
extra token优化, mtp layer重计算和tp支持 #2242
已合并
从refs/pull/2242/head合入到master
共 7 个文件变更+87-184
| @@ -21,16 +21,6 @@ from .blended_megatron_dataset_builder import need_to_build_dataset | |||
| 21 | logger = logging.getLogger(__name__) | 21 | logger = logging.getLogger(__name__) |
| 22 | 22 | ||
| 23 | 23 | ||
| 24 | -def gpt_dataset_init_wrapper(fn): | ||
| 25 | - | ||
| 26 | - def wrapper(self, *args, **kwargs): | ||
| 27 | - # Adapt to MTP | ||
| 28 | - _args = get_args() | ||
| 29 | - self.num_nextn_predict_layers = _args.num_nextn_predict_layers | ||
| 30 | - fn(self, *args, **kwargs) | ||
| 31 | - | ||
| 32 | - return wrapper | ||
| 33 | - | ||
| 34 | def gpt_dataset_getitem_wrapper(fn): | 24 | def gpt_dataset_getitem_wrapper(fn): |
| 35 | 25 | ||
| 36 | def wrapper(self, idx): | 26 | def wrapper(self, idx): |
| @@ -55,84 +45,6 @@ def gpt_dataset_getitem_wrapper(fn): | |||
| 55 | return wrapper | 45 | return wrapper |
| 56 | 46 | ||
| 57 | 47 | ||
| 58 | -def _query_document_sample_shuffle_indices( | ||
| 59 | - self, idx: int | ||
| 60 | -) -> Tuple[numpy.ndarray, numpy.ndarray]: | ||
| 61 | - """Get the text (token ids) and document ids for a given index | ||
| 62 | - | ||
| 63 | - Args: | ||
| 64 | - idx (int): The index into the dataset | ||
| 65 | - | ||
| 66 | - Returns: | ||
| 67 | - Tuple[numpy.ndarray, numpy.ndarray]: The text ids and document ids | ||
| 68 | - """ | ||
| 69 | - # Do the shuffle mapping | ||
| 70 | - idx = self.shuffle_index[idx] | ||
| 71 | - | ||
| 72 | - # Get the beginning and end documents and offsets | ||
| 73 | - doc_index_beg, doc_index_beg_offset = self.sample_index[idx] | ||
| 74 | - doc_index_end, doc_index_end_offset = self.sample_index[idx + 1] | ||
| 75 | - | ||
| 76 | - document_ids = [] | ||
| 77 | - sample_parts = [] | ||
| 78 | - | ||
| 79 | - # Sample spans a single document | ||
| 80 | - if doc_index_beg == doc_index_end: | ||
| 81 | - # Add the document id | ||
| 82 | - document_ids.append(self.document_index[doc_index_beg]) | ||
| 83 | - | ||
| 84 | - # Add the entire sample | ||
| 85 | - # Adapt to MTP | ||
| 86 | - sample_parts.append( | ||
| 87 | - self.dataset.get( | ||
| 88 | - self.document_index[doc_index_beg], | ||
| 89 | - offset=doc_index_beg_offset, | ||
| 90 | - length=doc_index_end_offset | ||
| 91 | - - doc_index_beg_offset | ||
| 92 | - + self.config.add_extra_token_to_sequence + self.num_nextn_predict_layers, | ||
| 93 | - ) | ||
| 94 | - ) | ||
| 95 | - | ||
| 96 | - # Sample spans multiple documents | ||
| 97 | - else: | ||
| 98 | - for i in range(doc_index_beg, doc_index_end + 1): | ||
| 99 | - # Add the document id | ||
| 100 | - document_ids.append(self.document_index[i]) | ||
| 101 | - | ||
| 102 | - # Add the sample part | ||
| 103 | - offset = 0 if i > doc_index_beg else doc_index_beg_offset | ||
| 104 | - # Adapt to MTP | ||
| 105 | - length = ( | ||
| 106 | - None | ||
| 107 | - if i < doc_index_end | ||
| 108 | - else doc_index_end_offset + self.config.add_extra_token_to_sequence + self.num_nextn_predict_layers | ||
| 109 | - ) | ||
| 110 | - sample_parts.append( | ||
| 111 | - self.dataset.get(self.document_index[i], offset=offset, length=length) | ||
| 112 | - ) | ||
| 113 | - assert len(document_ids) == len( | ||
| 114 | - sample_parts | ||
| 115 | - ), f"len(document_ids) ({len(document_ids)}) != len(sample_parts) ({len(sample_parts)})" | ||
| 116 | - | ||
| 117 | - length = sum(map(len, sample_parts)) | ||
| 118 | - | ||
| 119 | - # Pad the sample if necessary | ||
| 120 | - # Adapt to MTP | ||
| 121 | - if length < ( | ||
| 122 | - self.config.sequence_length + self.config.add_extra_token_to_sequence + self.num_nextn_predict_layers): | ||
| 123 | - sample_parts.append( | ||
| 124 | - [self._pad_token_id] | ||
| 125 | - * ( | ||
| 126 | - self.config.sequence_length + self.config.add_extra_token_to_sequence + | ||
| 127 | - self.num_nextn_predict_layers - length) | ||
| 128 | - ) | ||
| 129 | - | ||
| 130 | - return ( | ||
| 131 | - numpy.concatenate(sample_parts, dtype=numpy.int64), | ||
| 132 | - numpy.array(document_ids, dtype=numpy.int64), | ||
| 133 | - ) | ||
| 134 | - | ||
| 135 | - | ||
| 136 | def _build_document_sample_shuffle_indices( | 48 | def _build_document_sample_shuffle_indices( |
| 137 | self, | 49 | self, |
| 138 | ) -> Tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]: | 50 | ) -> Tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]: |
| @@ -206,16 +118,13 @@ def _build_document_sample_shuffle_indices( | |||
| 206 | separate_final_epoch = False | 118 | separate_final_epoch = False |
| 207 | else: | 119 | else: |
| 208 | # Get the number of samples for the last epoch | 120 | # Get the number of samples for the last epoch |
| 209 | - # Adapt to MTP | ||
| 210 | num_samples_sans_final_epoch = ( | 121 | num_samples_sans_final_epoch = ( |
| 211 | (num_epochs - 1) * num_tokens_per_epoch | 122 | (num_epochs - 1) * num_tokens_per_epoch |
| 212 | - self.config.add_extra_token_to_sequence | 123 | - self.config.add_extra_token_to_sequence |
| 213 | - - self.num_nextn_predict_layers | ||
| 214 | ) // sequence_length | 124 | ) // sequence_length |
| 215 | num_samples_from_final_epoch = self.num_samples - num_samples_sans_final_epoch | 125 | num_samples_from_final_epoch = self.num_samples - num_samples_sans_final_epoch |
| 216 | num_samples_per_epoch = ( | 126 | num_samples_per_epoch = ( |
| 217 | num_tokens_per_epoch - self.config.add_extra_token_to_sequence | 127 | num_tokens_per_epoch - self.config.add_extra_token_to_sequence |
| 218 | - - self.num_nextn_predict_layers | ||
| 219 | ) // sequence_length | 128 | ) // sequence_length |
| 220 | 129 | ||
| 221 | # num_samples_from_final_epoch should be non-negative | 130 | # num_samples_from_final_epoch should be non-negative |
| @@ -273,7 +182,6 @@ def _build_document_sample_shuffle_indices( | |||
| 273 | sequence_lengths_for_cpp = self.dataset.sequence_lengths.copy() | 182 | sequence_lengths_for_cpp = self.dataset.sequence_lengths.copy() |
| 274 | else: | 183 | else: |
| 275 | sequence_lengths_for_cpp = self.dataset.sequence_lengths | 184 | sequence_lengths_for_cpp = self.dataset.sequence_lengths |
| 276 | - # Adapt to MTP | ||
| 277 | sample_index = helpers.build_sample_idx( | 185 | sample_index = helpers.build_sample_idx( |
| 278 | sequence_lengths_for_cpp, | 186 | sequence_lengths_for_cpp, |
| 279 | document_index, | 187 | document_index, |
| @@ -281,7 +189,7 @@ def _build_document_sample_shuffle_indices( | |||
| 281 | num_epochs, | 189 | num_epochs, |
| 282 | num_tokens_per_epoch, | 190 | num_tokens_per_epoch, |
| 283 | drop_last_partial_sequence, | 191 | drop_last_partial_sequence, |
| 284 | - self.config.add_extra_token_to_sequence + self.num_nextn_predict_layers, | 192 | + self.config.add_extra_token_to_sequence, |
| 285 | ) | 193 | ) |
| 286 | 194 | ||
| 287 | if any(sample_index[:, 0] < 0): | 195 | if any(sample_index[:, 0] < 0): |
| @@ -368,39 +276,6 @@ def _build_document_sample_shuffle_indices( | |||
| 368 | return document_index, sample_index, shuffle_index | 276 | return document_index, sample_index, shuffle_index |
| 369 | 277 | ||
| 370 | 278 | ||
| 371 | -def _get_num_tokens_per_epoch(self) -> int: | ||
| 372 | - """Calculate the number of tokens in a single epoch | ||
| 373 | - | ||
| 374 | - Returns: | ||
| 375 | - int: The number of tokens in a single epoch | ||
| 376 | - """ | ||
| 377 | - return int(numpy.sum(self.dataset.sequence_lengths[self.indices])) + self.num_nextn_predict_layers | ||
| 378 | - | ||
| 379 | - | ||
| 380 | -def _get_num_epochs(self, num_tokens_per_epoch: int) -> int: | ||
| 381 | - """Calculate the number of epochs | ||
| 382 | - | ||
| 383 | - Args: | ||
| 384 | - num_tokens_per_epoch (int): The number of tokens in a single epoch | ||
| 385 | - | ||
| 386 | - Returns: | ||
| 387 | - int: The number of epochs | ||
| 388 | - """ | ||
| 389 | - num_epochs = 1 | ||
| 390 | - num_tokens = num_tokens_per_epoch | ||
| 391 | - if self.num_samples is None: | ||
| 392 | - return num_epochs | ||
| 393 | - else: | ||
| 394 | - # Adapt to MTP | ||
| 395 | - num_tokens_requested = ( | ||
| 396 | - self.num_samples * self.config.sequence_length | ||
| 397 | - ) + self.config.add_extra_token_to_sequence + self.num_nextn_predict_layers | ||
| 398 | - while num_tokens < num_tokens_requested: | ||
| 399 | - num_epochs += 1 | ||
| 400 | - num_tokens += num_tokens_per_epoch | ||
| 401 | - return num_epochs | ||
| 402 | - | ||
| 403 | - | ||
| 404 | def _get_ltor_masks_and_position_ids( | 279 | def _get_ltor_masks_and_position_ids( |
| 405 | data: torch.Tensor, | 280 | data: torch.Tensor, |
| 406 | eod_token: int, | 281 | eod_token: int, |
| @@ -432,7 +307,7 @@ def _get_ltor_masks_and_position_ids( | |||
| 432 | torch.Tensor: The position ID's of the token | 307 | torch.Tensor: The position ID's of the token |
| 433 | """ | 308 | """ |
| 434 | args = get_args() | 309 | args = get_args() |
| 435 | - seq_length = data.numel() - args.num_nextn_predict_layers | 310 | + seq_length = data.numel() |
| 436 | 311 | ||
| 437 | if create_attention_mask: | 312 | if create_attention_mask: |
| 438 | attention_mask = torch.tril( | 313 | attention_mask = torch.tril( |
| @@ -310,26 +310,35 @@ def inputs_slice(slice_num, input_ids, labels, position_ids, attention_mask): | |||
| 310 | [attention_mask], | 310 | [attention_mask], |
| 311 | ) | 311 | ) |
| 312 | 312 | ||
| 313 | - window_size = input_ids.shape[-1] - slice_num | ||
| 314 | return ( | 313 | return ( |
| 315 | - tensor_slide(input_ids, window_size), | 314 | + tensor_slide(input_ids, slice_num), |
| 316 | - tensor_slide(labels, window_size), | 315 | + tensor_slide(labels, slice_num), |
| 317 | - generate_nextn_position_ids(position_ids, window_size), | 316 | + generate_nextn_position_ids(position_ids, slice_num), |
| 318 | # not compatible with ppo attn_mask | 317 | # not compatible with ppo attn_mask |
| 319 | - tensor_slide(attention_mask, window_size, dims=[-2, -1]), | 318 | + tensor_slide(attention_mask, slice_num, dims=[-2, -1]), |
| 320 | ) | 319 | ) |
| 321 | 320 | ||
| 322 | 321 | ||
| 323 | -def generate_nextn_position_ids(tensor, window_size): | 322 | +def generate_nextn_position_ids(tensor, slice_num): |
| 324 | - slides = tensor_slide(tensor, window_size) | 323 | + slides = tensor_slide(tensor, slice_num) |
| 324 | + if slides[0] is None: | ||
| 325 | + return slides | ||
| 325 | 326 | ||
| 326 | for idx in range(1, len(slides)): | 327 | for idx in range(1, len(slides)): |
| 327 | - for i in range(slides[idx].size(0)): | 328 | + slides[idx] = regenerate_position_ids(slides[idx], idx) |
| 328 | - row = slides[idx][i] | ||
| 329 | - zero_mask = (row == 0) | ||
| 330 | - if zero_mask.any(): | ||
| 331 | - first_zero_idx = torch.argmax(zero_mask.int()).item() | ||
| 332 | - slides[idx][i, :first_zero_idx] = torch.arange(first_zero_idx) | ||
| 333 | - else: | ||
| 334 | - slides[idx] = slides[idx] - idx | ||
| 335 | return slides | 329 | return slides |
| 330 | + | ||
| 331 | + | ||
| 332 | +def regenerate_position_ids(tensor, offset): | ||
| 333 | + if tensor is None: | ||
| 334 | + return None | ||
| 335 | + tensor = tensor.clone() | ||
| 336 | + for i in range(tensor.size(0)): | ||
| 337 | + row = tensor[i] | ||
| 338 | + zero_mask = (row == 0) | ||
| 339 | + if zero_mask.any(): | ||
| 340 | + first_zero_idx = torch.argmax(zero_mask.int()).item() | ||
| 341 | + tensor[i, :first_zero_idx] = torch.arange(first_zero_idx) | ||
| 342 | + else: | ||
| 343 | + tensor = tensor - offset | ||
| 344 | + return tensor | ||
| @@ -518,24 +518,11 @@ class CoreAdaptation(MegatronAdaptationABC): | |||
| 518 | finalize_wrapper) | 518 | finalize_wrapper) |
| 519 | # MTP need extra token | 519 | # MTP need extra token |
| 520 | from ..core.datasets.gpt_dataset import ( | 520 | from ..core.datasets.gpt_dataset import ( |
| 521 | - gpt_dataset_init_wrapper, | ||
| 522 | gpt_dataset_getitem_wrapper, | 521 | gpt_dataset_getitem_wrapper, |
| 523 | - _query_document_sample_shuffle_indices, | ||
| 524 | - _get_num_tokens_per_epoch, | ||
| 525 | - _get_num_epochs, | ||
| 526 | _get_ltor_masks_and_position_ids | 522 | _get_ltor_masks_and_position_ids |
| 527 | ) | 523 | ) |
| 528 | - MegatronAdaptation.register('megatron.core.datasets.gpt_dataset.GPTDataset.__init__', | ||
| 529 | - gpt_dataset_init_wrapper) | ||
| 530 | MegatronAdaptation.register('megatron.core.datasets.gpt_dataset.GPTDataset.__getitem__', | 524 | MegatronAdaptation.register('megatron.core.datasets.gpt_dataset.GPTDataset.__getitem__', |
| 531 | gpt_dataset_getitem_wrapper) | 525 | gpt_dataset_getitem_wrapper) |
| 532 | - MegatronAdaptation.register( | ||
| 533 | - 'megatron.core.datasets.gpt_dataset.GPTDataset._query_document_sample_shuffle_indices', | ||
| 534 | - _query_document_sample_shuffle_indices) | ||
| 535 | - MegatronAdaptation.register('megatron.core.datasets.gpt_dataset.GPTDataset._get_num_tokens_per_epoch', | ||
| 536 | - _get_num_tokens_per_epoch) | ||
| 537 | - MegatronAdaptation.register('megatron.core.datasets.gpt_dataset.GPTDataset._get_num_epochs', | ||
| 538 | - _get_num_epochs) | ||
| 539 | MegatronAdaptation.register('megatron.core.datasets.gpt_dataset._get_ltor_masks_and_position_ids', | 526 | MegatronAdaptation.register('megatron.core.datasets.gpt_dataset._get_ltor_masks_and_position_ids', |
| 540 | _get_ltor_masks_and_position_ids) | 527 | _get_ltor_masks_and_position_ids) |
| 541 | 528 | ||
| @@ -74,6 +74,7 @@ class MultiTokenPredication(MegatronModule): | |||
| 74 | # share with main model | 74 | # share with main model |
| 75 | self.share_mtp_embedding_and_output_weight = share_mtp_embedding_and_output_weight | 75 | self.share_mtp_embedding_and_output_weight = share_mtp_embedding_and_output_weight |
| 76 | self.recompute_layer_norm = args.recompute_mtp_norm | 76 | self.recompute_layer_norm = args.recompute_mtp_norm |
| 77 | + self.recompute_mtp_layer = args.recompute_mtp_layer | ||
| 77 | 78 | ||
| 78 | self.embedding = LanguageModelEmbedding( | 79 | self.embedding = LanguageModelEmbedding( |
| 79 | config=self.config, | 80 | config=self.config, |
| @@ -243,14 +244,32 @@ class MultiTokenPredication(MegatronModule): | |||
| 243 | hidden_states.register_hook(self.hnorm_ckpt.recompute) | 244 | hidden_states.register_hook(self.hnorm_ckpt.recompute) |
| 244 | # hidden_states -> [s, b, h] | 245 | # hidden_states -> [s, b, h] |
| 245 | hidden_states, _ = self.eh_proj(hidden_states) | 246 | hidden_states, _ = self.eh_proj(hidden_states) |
| 246 | - hidden_states, _ = self.transformer_layer( | 247 | + |
| 247 | - hidden_states=hidden_states, | 248 | + if self.config.tensor_model_parallel_size > 1: |
| 248 | - attention_mask=attention_mask, | 249 | + hidden_states = tensor_parallel.gather_from_tensor_model_parallel_region(hidden_states) |
| 249 | - rotary_pos_emb=rotary_pos_emb, | 250 | + if self.config.sequence_parallel: |
| 250 | - inference_params=inference_params, | 251 | + hidden_states = tensor_parallel.scatter_to_sequence_parallel_region(hidden_states) |
| 251 | - packed_seq_params=packed_seq_params, | 252 | + if self.recompute_mtp_layer: |
| 252 | - **(extra_block_kwargs or {}), | 253 | + hidden_states, context = tensor_parallel.checkpoint( |
| 253 | - ) | 254 | + self.transformer_layer, |
| 255 | + self.config.distribute_saved_activations, | ||
| 256 | + hidden_states, | ||
| 257 | + attention_mask, | ||
| 258 | + None, | ||
| 259 | + None, | ||
| 260 | + rotary_pos_emb, | ||
| 261 | + inference_params, | ||
| 262 | + packed_seq_params, | ||
| 263 | + ) | ||
| 264 | + else: | ||
| 265 | + hidden_states, _ = self.transformer_layer( | ||
| 266 | + hidden_states=hidden_states, | ||
| 267 | + attention_mask=attention_mask, | ||
| 268 | + rotary_pos_emb=rotary_pos_emb, | ||
| 269 | + inference_params=inference_params, | ||
| 270 | + packed_seq_params=packed_seq_params, | ||
| 271 | + **(extra_block_kwargs or {}), | ||
| 272 | + ) | ||
| 254 | 273 | ||
| 255 | # Final layer norm. | 274 | # Final layer norm. |
| 256 | if self.final_layernorm is not None: | 275 | if self.final_layernorm is not None: |
| @@ -156,6 +156,8 @@ def _add_mtp_args(parser): | |||
| 156 | group.add_argument('--mtp-loss-scale', type=float, default=0.3, help='Multi-Token prediction loss scale') | 156 | group.add_argument('--mtp-loss-scale', type=float, default=0.3, help='Multi-Token prediction loss scale') |
| 157 | group.add_argument('--recompute-mtp-norm', action='store_true', default=False, | 157 | group.add_argument('--recompute-mtp-norm', action='store_true', default=False, |
| 158 | help='Multi-Token prediction recompute norm') | 158 | help='Multi-Token prediction recompute norm') |
| 159 | + group.add_argument('--recompute-mtp-layer', action='store_true', default=False, | ||
| 160 | + help='Multi-Token prediction recompute layer') | ||
| 159 | group.add_argument('--share-mtp-embedding-and-output-weight', action='store_true', default=False, | 161 | group.add_argument('--share-mtp-embedding-and-output-weight', action='store_true', default=False, |
| 160 | help='Main model share embedding and output weight with mtp layer.') | 162 | help='Main model share embedding and output weight with mtp layer.') |
| 161 | return parser | 163 | return parser |
| @@ -294,10 +294,12 @@ def get_batch_on_this_tp_rank(data_iterator): | |||
| 294 | _broadcast(batch['position_ids']) | 294 | _broadcast(batch['position_ids']) |
| 295 | 295 | ||
| 296 | elif mpu.is_pipeline_last_stage(): | 296 | elif mpu.is_pipeline_last_stage(): |
| 297 | + if args.num_nextn_predict_layers: | ||
| 298 | + _broadcast(batch['tokens']) | ||
| 297 | _broadcast(batch['labels']) | 299 | _broadcast(batch['labels']) |
| 298 | _broadcast(batch['loss_mask']) | 300 | _broadcast(batch['loss_mask']) |
| 299 | _broadcast(batch['attention_mask']) | 301 | _broadcast(batch['attention_mask']) |
| 300 | - if args.reset_position_ids: | 302 | + if args.reset_position_ids or args.num_nextn_predict_layers: |
| 301 | _broadcast(batch['position_ids']) | 303 | _broadcast(batch['position_ids']) |
| 302 | else: | 304 | else: |
| 303 | _broadcast(batch['attention_mask']) | 305 | _broadcast(batch['attention_mask']) |
| @@ -306,20 +308,25 @@ def get_batch_on_this_tp_rank(data_iterator): | |||
| 306 | 308 | ||
| 307 | else: | 309 | else: |
| 308 | 310 | ||
| 309 | - tokens = torch.empty((args.micro_batch_size, args.seq_length), dtype=torch.int64, | 311 | + tokens = torch.empty((args.micro_batch_size, args.seq_length + args.num_nextn_predict_layers), |
| 312 | + dtype=torch.int64, | ||
| 310 | device=torch.cuda.current_device()) | 313 | device=torch.cuda.current_device()) |
| 311 | - labels = torch.empty((args.micro_batch_size, args.seq_length), dtype=torch.int64, | 314 | + labels = torch.empty((args.micro_batch_size, args.seq_length + args.num_nextn_predict_layers), |
| 315 | + dtype=torch.int64, | ||
| 312 | device=torch.cuda.current_device()) | 316 | device=torch.cuda.current_device()) |
| 313 | - loss_mask = torch.empty((args.micro_batch_size, args.seq_length), dtype=torch.float32, | 317 | + loss_mask = torch.empty((args.micro_batch_size, args.seq_length + args.num_nextn_predict_layers), |
| 318 | + dtype=torch.float32, | ||
| 314 | device=torch.cuda.current_device()) | 319 | device=torch.cuda.current_device()) |
| 315 | if args.create_attention_mask_in_dataloader: | 320 | if args.create_attention_mask_in_dataloader: |
| 316 | attention_mask = torch.empty( | 321 | attention_mask = torch.empty( |
| 317 | - (args.micro_batch_size, 1, args.seq_length, args.seq_length), dtype=torch.bool, | 322 | + (args.micro_batch_size, 1, args.seq_length + args.num_nextn_predict_layers, |
| 323 | + args.seq_length + args.num_nextn_predict_layers), dtype=torch.bool, | ||
| 318 | device=torch.cuda.current_device() | 324 | device=torch.cuda.current_device() |
| 319 | ) | 325 | ) |
| 320 | else: | 326 | else: |
| 321 | attention_mask = None | 327 | attention_mask = None |
| 322 | - position_ids = torch.empty((args.micro_batch_size, args.seq_length), dtype=torch.int64, | 328 | + position_ids = torch.empty((args.micro_batch_size, args.seq_length + args.num_nextn_predict_layers), |
| 329 | + dtype=torch.int64, | ||
| 323 | device=torch.cuda.current_device()) | 330 | device=torch.cuda.current_device()) |
| 324 | 331 | ||
| 325 | if args.pipeline_model_parallel_size == 1: | 332 | if args.pipeline_model_parallel_size == 1: |
| @@ -337,11 +344,14 @@ def get_batch_on_this_tp_rank(data_iterator): | |||
| 337 | _broadcast(position_ids) | 344 | _broadcast(position_ids) |
| 338 | 345 | ||
| 339 | elif mpu.is_pipeline_last_stage(): | 346 | elif mpu.is_pipeline_last_stage(): |
| 340 | - tokens = None | 347 | + if args.num_nextn_predict_layers: |
| 348 | + _broadcast(tokens) | ||
| 349 | + else: | ||
| 350 | + tokens = None | ||
| 341 | _broadcast(labels) | 351 | _broadcast(labels) |
| 342 | _broadcast(loss_mask) | 352 | _broadcast(loss_mask) |
| 343 | _broadcast(attention_mask) | 353 | _broadcast(attention_mask) |
| 344 | - if args.reset_position_ids: | 354 | + if args.reset_position_ids or args.num_nextn_predict_layers: |
| 345 | _broadcast(position_ids) | 355 | _broadcast(position_ids) |
| 346 | else: | 356 | else: |
| 347 | position_ids = None | 357 | position_ids = None |
| @@ -511,26 +521,27 @@ def _get_batch_on_this_cp_rank_in_megatron_cp_general(batch): | |||
| 511 | 521 | ||
| 512 | def tensor_slide( | 522 | def tensor_slide( |
| 513 | tensor: Optional[torch.Tensor], | 523 | tensor: Optional[torch.Tensor], |
| 514 | - window_size: int = None, | 524 | + slice_num: int, |
| 515 | dims: Union[int, List[int]] = -1, | 525 | dims: Union[int, List[int]] = -1, |
| 516 | - step: int = 1 | 526 | + step: int = 1, |
| 527 | + return_first=False, | ||
| 517 | ) -> List[Union[torch.Tensor, None]]: | 528 | ) -> List[Union[torch.Tensor, None]]: |
| 518 | - """slide window slice for n-D tensor""" | 529 | + """通用滑动窗口函数，支持任意维度""" |
| 519 | if tensor is None: | 530 | if tensor is None: |
| 520 | # return `List[None]` to avoid NoneType Error | 531 | # return `List[None]` to avoid NoneType Error |
| 521 | - return [None] | 532 | + return [None] * (slice_num + 1) |
| 522 | - | 533 | + if slice_num == 0: |
| 523 | - if window_size is None: | ||
| 524 | - window_size = tensor.shape[-1] | ||
| 525 | - if window_size == tensor.shape[-1]: | ||
| 526 | return [tensor] | 534 | return [tensor] |
| 527 | - | 535 | + window_size = tensor.shape[-1] - slice_num |
| 528 | dims = [dims] if isinstance(dims, int) else sorted(dims, reverse=True) | 536 | dims = [dims] if isinstance(dims, int) else sorted(dims, reverse=True) |
| 529 | 537 | ||
| 538 | + # 连续多维度滑动 | ||
| 530 | slices = [] | 539 | slices = [] |
| 531 | for i in range(0, tensor.size(dims[-1]) - window_size + 1, step): | 540 | for i in range(0, tensor.size(dims[-1]) - window_size + 1, step): |
| 532 | slice_obj = [slice(None)] * tensor.dim() | 541 | slice_obj = [slice(None)] * tensor.dim() |
| 533 | for dim in dims: | 542 | for dim in dims: |
| 534 | slice_obj[dim] = slice(i, i + window_size) | 543 | slice_obj[dim] = slice(i, i + window_size) |
| 535 | - slices.append(tensor[tuple(slice_obj)].clone()) | 544 | + slices.append(tensor[tuple(slice_obj)]) |
| 545 | + if return_first: | ||
| 546 | + return slices | ||
| 536 | return slices | 547 | return slices |
| @@ -124,14 +124,14 @@ def loss_func(loss_mask: torch.Tensor, output_tensor: torch.Tensor): | |||
| 124 | args = get_args() | 124 | args = get_args() |
| 125 | 125 | ||
| 126 | losses = output_tensor.float() | 126 | losses = output_tensor.float() |
| 127 | + if args.num_nextn_predict_layers > 0: | ||
| 128 | + loss_mask = tensor_slide(loss_mask, args.num_nextn_predict_layers, return_first=True)[0] | ||
| 127 | loss_mask = loss_mask.view(-1).float() | 129 | loss_mask = loss_mask.view(-1).float() |
| 128 | if args.context_parallel_size > 1: | 130 | if args.context_parallel_size > 1: |
| 129 | loss = torch.cat([torch.sum(losses.view(-1) * loss_mask).view(1), loss_mask.sum().view(1)]) | 131 | loss = torch.cat([torch.sum(losses.view(-1) * loss_mask).view(1), loss_mask.sum().view(1)]) |
| 130 | torch.distributed.all_reduce(loss, group=mpu.get_context_parallel_group()) | 132 | torch.distributed.all_reduce(loss, group=mpu.get_context_parallel_group()) |
| 131 | loss = loss[0] / loss[1] | 133 | loss = loss[0] / loss[1] |
| 132 | else: | 134 | else: |
| 133 | - if args.num_nextn_predict_layers > 0: | ||
| 134 | - loss_mask = tensor_slide(loss_mask, loss_mask.shape[-1] - args.num_nextn_predict_layers)[0] | ||
| 135 | loss = torch.sum(losses.view(-1) * loss_mask) / loss_mask.sum() | 135 | loss = torch.sum(losses.view(-1) * loss_mask) / loss_mask.sum() |
| 136 | 136 | ||
| 137 | # Check individual rank losses are not NaN prior to DP all-reduce. | 137 | # Check individual rank losses are not NaN prior to DP all-reduce. |
| @@ -178,7 +178,7 @@ def core_gpt_dataset_config_from_args(args): | |||
| 178 | 178 | ||
| 179 | return GPTDatasetConfig( | 179 | return GPTDatasetConfig( |
| 180 | random_seed=args.seed, | 180 | random_seed=args.seed, |
| 181 | - sequence_length=args.seq_length, | 181 | + sequence_length=args.seq_length + args.num_nextn_predict_layers, |
| 182 | blend=get_blend_from_list(args.data_path), | 182 | blend=get_blend_from_list(args.data_path), |
| 183 | blend_per_split=[ | 183 | blend_per_split=[ |
| 184 | get_blend_from_list(args.train_data_path), | 184 | get_blend_from_list(args.train_data_path), |