已合并
[pytorch][feature]DSA indexer loss #3631
iansheng创建于 2025年11月4日
[pytorch][feature]DSA indexer loss #3631
已合并
共 5 个文件变更+604-12
| @@ -9,15 +9,17 @@ class DSAIndexerFeature(MindSpeedFeature): | |||
| 9 | group = parser.add_argument_group(title=self.feature_name) | 9 | group = parser.add_argument_group(title=self.feature_name) |
| 10 | 10 | ||
| 11 | group.add_argument('--enable-dsa-indexer', action='store_true', default=False, | 11 | group.add_argument('--enable-dsa-indexer', action='store_true', default=False, |
| 12 | help='add dsa_indexer module in MLA.') | 12 | help='add dsa_indexer module in MLA.') |
| 13 | group.add_argument('--index-n-heads', type=int, default=64, | 13 | group.add_argument('--index-n-heads', type=int, default=64, |
| 14 | help='dimension for index head number.') | 14 | help='dimension for index head number.') |
| 15 | group.add_argument('--index-head-dim', type=int, default=128, | 15 | group.add_argument('--index-head-dim', type=int, default=128, |
| 16 | help='dimension for index head dim.') | 16 | help='dimension for index head dim.') |
| 17 | group.add_argument('--index-topk', type=int, default=2048, | 17 | group.add_argument('--index-topk', type=int, default=2048, |
| 18 | help='top-k for index head') | 18 | help='top-k for index head') |
| 19 | group.add_argument('--scale-fmt', type=str, default=None, | 19 | group.add_argument('--scale-fmt', type=str, default=None, |
| 20 | help='format for quantization scale.') | 20 | help='format for quantization scale.') |
| 21 | group.add_argument('--indexer-loss-coeff', type=float, default=1.0, | ||
| 22 | help='Indexer loss coeff.') | ||
| 21 | 23 | ||
| 22 | def validate_args(self, args): | 24 | def validate_args(self, args): |
| 23 | if args.enable_dsa_indexer: | 25 | if args.enable_dsa_indexer: |
| @@ -30,4 +32,7 @@ class DSAIndexerFeature(MindSpeedFeature): | |||
| 30 | if args.enable_dsa_indexer: | 32 | if args.enable_dsa_indexer: |
| 31 | from mindspeed_llm.tasks.models.transformer.dsa_indexer import fp16module_init_wrapper | 33 | from mindspeed_llm.tasks.models.transformer.dsa_indexer import fp16module_init_wrapper |
| 32 | patch_manager.register_patch('megatron.core.transformer.module.Float16Module.__init__', | 34 | patch_manager.register_patch('megatron.core.transformer.module.Float16Module.__init__', |
| 33 | fp16module_init_wrapper) | 35 | fp16module_init_wrapper) |
| 36 | from mindspeed_llm.tasks.models.transformer.dsa_indexer import forward_step_dsa_wrapper | ||
| 37 | patch_manager.register_patch('megatron.core.pipeline_parallel.schedules.forward_step', | ||
| 38 | forward_step_dsa_wrapper) | ||
| @@ -1,3 +1,4 @@ | |||
| 1 | import contextlib | ||
| 1 | import math | 2 | import math |
| 2 | from dataclasses import dataclass | 3 | from dataclasses import dataclass |
| 3 | from typing import List, Tuple, Union, Optional | 4 | from typing import List, Tuple, Union, Optional |
| @@ -7,6 +8,12 @@ import torch.nn.functional as F | |||
| 7 | from einops import rearrange | 8 | from einops import rearrange |
| 8 | from functools import wraps | 9 | from functools import wraps |
| 9 | 10 | ||
| 11 | from megatron.core import parallel_state | ||
| 12 | from megatron.core.enums import ModelType | ||
| 13 | from megatron.core.pipeline_parallel.schedules import set_current_microbatch | ||
| 14 | from megatron.core.transformer.moe.moe_utils import MoEAuxLossAutoScaler | ||
| 15 | from megatron.core.transformer.multi_token_prediction import MTPLossAutoScaler | ||
| 16 | from megatron.core.utils import get_attr_wrapped_model, get_model_type | ||
| 10 | from megatron.training import get_args | 17 | from megatron.training import get_args |
| 11 | from megatron.legacy.model import RMSNorm | 18 | from megatron.legacy.model import RMSNorm |
| 12 | from megatron.core.transformer.identity_op import IdentityOp | 19 | from megatron.core.transformer.identity_op import IdentityOp |
| @@ -287,3 +294,273 @@ class DSAIndexer(MegatronModule): | |||
| 287 | args.sparse_mode = 0 | 294 | args.sparse_mode = 0 |
| 288 | 295 | ||
| 289 | return topk_score, topk_indices, attention_mask | 296 | return topk_score, topk_indices, attention_mask |
| 297 | |||
| 298 | |||
| 299 | class DSAIndexerLossAutoScaler(torch.autograd.Function): | ||
| 300 | """An AutoScaler that triggers the backward pass and scales the grad for DSA indexer loss.""" | ||
| 301 | |||
| 302 | main_loss_backward_scale: torch.Tensor = None | ||
| 303 | |||
| 304 | |||
| 305 | def forward(ctx, output: torch.Tensor, aux_loss: torch.Tensor): | ||
| 306 | """Preserve the indexer_loss by storing it in the context to avoid garbage collection. | ||
| 307 | |||
| 308 | Args: | ||
| 309 | output (torch.Tensor): The output tensor. | ||
| 310 | aux_loss (torch.Tensor): The indexer loss tensor. | ||
| 311 | |||
| 312 | Returns: | ||
| 313 | torch.Tensor: The output tensor. | ||
| 314 | """ | ||
| 315 | ctx.save_for_backward(aux_loss) | ||
| 316 | return output | ||
| 317 | |||
| 318 | |||
| 319 | def backward(ctx, grad_output: torch.Tensor): | ||
| 320 | """Compute and scale the gradient for indexer loss. | ||
| 321 | |||
| 322 | Args: | ||
| 323 | grad_output (torch.Tensor): The gradient of the output. | ||
| 324 | |||
| 325 | Returns: | ||
| 326 | Tuple[torch.Tensor, torch.Tensor]: The gradient of the output, scaled indexer loss | ||
| 327 | gradient. | ||
| 328 | """ | ||
| 329 | (loss,) = ctx.saved_tensors | ||
| 330 | if DSAIndexerLossAutoScaler.main_loss_backward_scale is None: | ||
| 331 | DSAIndexerLossAutoScaler.main_loss_backward_scale = torch.tensor( | ||
| 332 | 1.0, device=loss.device | ||
| 333 | ) | ||
| 334 | dsa_indexer_loss_backward_scale = DSAIndexerLossAutoScaler.main_loss_backward_scale | ||
| 335 | scaled_dsa_indexer_loss_grad = torch.ones_like(loss) * dsa_indexer_loss_backward_scale | ||
| 336 | return grad_output, scaled_dsa_indexer_loss_grad | ||
| 337 | |||
| 338 | |||
| 339 | def set_loss_scale(scale: torch.Tensor): | ||
| 340 | """set the scale of the indexer loss. | ||
| 341 | |||
| 342 | Args: | ||
| 343 | scale (torch.Tensor): The scale value to set. Please ensure that the scale passed in | ||
| 344 | matches the scale of the main_loss. | ||
| 345 | """ | ||
| 346 | if DSAIndexerLossAutoScaler.main_loss_backward_scale is None: | ||
| 347 | DSAIndexerLossAutoScaler.main_loss_backward_scale = scale | ||
| 348 | else: | ||
| 349 | DSAIndexerLossAutoScaler.main_loss_backward_scale.copy_(scale) | ||
| 350 | |||
| 351 | |||
| 352 | def forward_step_dsa_wrapper(fn): | ||
| 353 | """Forward step for passed-in model. Patch for DSA indexer loss. | ||
| 354 | """ | ||
| 355 | |||
| 356 | |||
| 357 | def wrapper( | ||
| 358 | forward_step_func, | ||
| 359 | data_iterator, | ||
| 360 | model, | ||
| 361 | num_microbatches, | ||
| 362 | input_tensor, | ||
| 363 | forward_data_store, | ||
| 364 | config, | ||
| 365 | collect_non_loss_data=False, | ||
| 366 | checkpoint_activations_microbatch=None, | ||
| 367 | is_first_microbatch=False, | ||
| 368 | current_microbatch=None, | ||
| 369 | encoder_decoder_xattn=False, | ||
| 370 | ): | ||
| 371 | output_tensor, num_tokens = fn( | ||
| 372 | forward_step_func, | ||
| 373 | data_iterator, | ||
| 374 | model, | ||
| 375 | num_microbatches, | ||
| 376 | input_tensor, | ||
| 377 | forward_data_store, | ||
| 378 | config, | ||
| 379 | collect_non_loss_data=collect_non_loss_data, | ||
| 380 | checkpoint_activations_microbatch=checkpoint_activations_microbatch, | ||
| 381 | is_first_microbatch=is_first_microbatch, | ||
| 382 | current_microbatch=current_microbatch, | ||
| 383 | encoder_decoder_xattn=encoder_decoder_xattn, | ||
| 384 | ) | ||
| 385 | if not isinstance(output_tensor, list): | ||
| 386 | output_tensor_device = output_tensor.device | ||
| 387 | else: | ||
| 388 | output_tensor_device = output_tensor[0].device | ||
| 389 | # Set the loss scale for DSA indexer loss. | ||
| 390 | global_args = get_args() | ||
| 391 | if global_args.enable_dsa_indexer: | ||
| 392 | # Calculate the loss scale based on the grad_scale_func if available, else default to 1. | ||
| 393 | loss_scale = ( | ||
| 394 | config.grad_scale_func(torch.ones(1, device=output_tensor_device)) | ||
| 395 | if config.grad_scale_func is not None | ||
| 396 | else torch.ones(1, device=output_tensor_device) | ||
| 397 | ) | ||
| 398 | # Set the loss scale | ||
| 399 | if config.calculate_per_token_loss: | ||
| 400 | DSAIndexerLossAutoScaler.set_loss_scale(loss_scale) | ||
| 401 | else: | ||
| 402 | DSAIndexerLossAutoScaler.set_loss_scale(loss_scale / num_microbatches) | ||
| 403 | return output_tensor, num_tokens | ||
| 404 | |||
| 405 | return wrapper | ||
| 406 | |||
| 407 | |||
| 408 | class DSAIndexerLossLoggingHelper: | ||
| 409 | """Helper class for logging DSAIndexer losses.""" | ||
| 410 | |||
| 411 | tracker = {} | ||
| 412 | |||
| 413 | |||
| 414 | def save_loss_to_tracker( | ||
| 415 | loss: torch.Tensor, | ||
| 416 | layer_number: int, | ||
| 417 | num_layers: int, | ||
| 418 | reduce_group: torch.distributed.ProcessGroup = None, | ||
| 419 | avg_group: torch.distributed.ProcessGroup = None, | ||
| 420 | ): | ||
| 421 | """Save the DSA indexer loss for logging. | ||
| 422 | Args: | ||
| 423 | loss (torch.Tensor): The loss tensor. | ||
| 424 | layer_number (int): Layer index of the loss. | ||
| 425 | num_layers (int): The number of total layers. | ||
| 426 | reduce_group (torch.distributed.ProcessGroup): The group for reducing the loss. | ||
| 427 | mean_group (torch.distributed.ProcessGroup): The group for averaging the loss. | ||
| 428 | """ | ||
| 429 | # Skip DSA indexer loss logging if layer_number is None. | ||
| 430 | if layer_number is None: | ||
| 431 | return | ||
| 432 | |||
| 433 | tracker = DSAIndexerLossLoggingHelper.tracker | ||
| 434 | if "values" not in tracker: | ||
| 435 | tracker["values"] = torch.zeros(num_layers, device=loss.device) | ||
| 436 | tracker["values"][layer_number - 1] += loss.detach() | ||
| 437 | tracker["reduce_group"] = reduce_group | ||
| 438 | tracker["avg_group"] = avg_group | ||
| 439 | |||
| 440 | |||
| 441 | def clean_loss_in_tracker(): | ||
| 442 | """Clear the DSA indexer losses.""" | ||
| 443 | tracker = DSAIndexerLossLoggingHelper.tracker | ||
| 444 | tracker["values"].zero_() | ||
| 445 | tracker["reduce_group"] = None | ||
| 446 | tracker["avg_group"] = None | ||
| 447 | |||
| 448 | |||
| 449 | def reduce_loss_in_tracker(): | ||
| 450 | """Collect and reduce the DSA indexer losses across ranks.""" | ||
| 451 | tracker = DSAIndexerLossLoggingHelper.tracker | ||
| 452 | if "values" not in tracker: | ||
| 453 | return | ||
| 454 | values = tracker["values"] | ||
| 455 | # Collect DSA indexer losses across PP. | ||
| 456 | torch.distributed.all_reduce( | ||
| 457 | values, group=parallel_state.get_pipeline_model_parallel_group() | ||
| 458 | ) | ||
| 459 | # Reduce DSA indexer losses across ranks. | ||
| 460 | if tracker.get('reduce_group') is not None: | ||
| 461 | torch.distributed.all_reduce(values, group=tracker.get('reduce_group')) | ||
| 462 | if tracker.get('avg_group') is not None: | ||
| 463 | torch.distributed.all_reduce( | ||
| 464 | values, group=tracker['avg_group'], op=torch.distributed.ReduceOp.AVG | ||
| 465 | ) | ||
| 466 | |||
| 467 | |||
| 468 | def track_das_indexer_metrics(loss_scale, iteration, writer, wandb_writer=None, total_loss_dict=None): | ||
| 469 | """Track the DSA Indexer metrics for logging.""" | ||
| 470 | DSAIndexerLossLoggingHelper.reduce_loss_in_tracker() | ||
| 471 | tracker = DSAIndexerLossLoggingHelper.tracker | ||
| 472 | if "values" not in tracker: | ||
| 473 | return | ||
| 474 | das_indexer_losses = tracker["values"] * loss_scale | ||
| 475 | das_indexer_num_layers = das_indexer_losses.shape[0] | ||
| 476 | loss = das_indexer_losses.sum() / das_indexer_num_layers | ||
| 477 | name = "dsa_indexer_loss" | ||
| 478 | if total_loss_dict is not None: | ||
| 479 | total_loss_dict[name] = loss | ||
| 480 | if writer is not None: | ||
| 481 | writer.add_scalar(name, loss, iteration) | ||
| 482 | if wandb_writer is not None: | ||
| 483 | wandb_writer.log({f"{name}": loss}, iteration) | ||
| 484 | |||
| 485 | DSAIndexerLossLoggingHelper.clean_loss_in_tracker() | ||
| 486 | |||
| 487 | def compute_dsa_indexer_loss( | ||
| 488 | main_attn_dist, | ||
| 489 | index_score, | ||
| 490 | topk_indices, | ||
| 491 | loss_scale, | ||
| 492 | ): | ||
| 493 | """Compute dsa indexer loss at sparse training stage | ||
| 494 | Reference: https://github.com/deepseek-ai/DeepSeek-V3.2-Exp/blob/main/DeepSeek_V3_2.pdf | ||
| 495 | Args: | ||
| 496 | main_attn_dist: Q dist | ||
| 497 | index_score: P dist | ||
| 498 | topk_indices: Selected top-K indices for sparse phase | ||
| 499 | loss_scale: Dsa indexer loss scale | ||
| 500 | """ | ||
| 501 | index_score = F.softmax(index_score, dim=-1, dtype=torch.float32) | ||
| 502 | # considering only the selected token | ||
| 503 | selected_main_attn_dist = torch.gather(main_attn_dist, dim=-1, index=topk_indices) | ||
| 504 | selected_main_attn_dist = F.normalize(selected_main_attn_dist, p=1, dim=-1) | ||
| 505 | loss = F.kl_div((index_score + 1e-10).log(), | ||
| 506 | selected_main_attn_dist + 1e-10, | ||
| 507 | reduction='none', | ||
| 508 | ).sum(dim=-1).mean() | ||
| 509 | loss *= loss_scale | ||
| 510 | |||
| 511 | return loss | ||
| 512 | |||
| 513 | |||
| 514 | def get_attn_scores( | ||
| 515 | query, | ||
| 516 | key, | ||
| 517 | attention_mask, | ||
| 518 | num_attn_head_per_group, | ||
| 519 | attn_scale, | ||
| 520 | ): | ||
| 521 | """aggregate the main attention scores""" | ||
| 522 | if num_attn_head_per_group > 1: | ||
| 523 | key = key.repeat_interleave( | ||
| 524 | num_attn_head_per_group, dim=2 | ||
| 525 | ) | ||
| 526 | |||
| 527 | # [b, np, sq, sk] | ||
| 528 | output_size = (query.size(1), query.size(2), query.size(0), key.size(0)) | ||
| 529 | |||
| 530 | # [sq, b, np, hn] -> [sq, b * np, hn] | ||
| 531 | # This will be a simple view when doing normal attention, but in group query attention | ||
| 532 | # the key and value tensors are repeated to match the queries so you can't use | ||
| 533 | # simple strides to extract the queries. | ||
| 534 | query = query.reshape(output_size[2], output_size[0] * output_size[1], -1) | ||
| 535 | # [sk, b, np, hn] -> [sk, b * np, hn] | ||
| 536 | key = key.view(output_size[3], output_size[0] * output_size[1], -1) | ||
| 537 | |||
| 538 | # preallocting input tensor: [b * np, sq, sk] | ||
| 539 | matmul_input_buffer = parallel_state.get_global_memory_buffer().get_tensor( | ||
| 540 | (output_size[0] * output_size[1], output_size[2], output_size[3]), query.dtype, "mpu" | ||
| 541 | ) | ||
| 542 | |||
| 543 | # Raw attention scores. [b * np, sq, sk] | ||
| 544 | matmul_result = torch.baddbmm( | ||
| 545 | matmul_input_buffer, | ||
| 546 | query.transpose(0, 1), # [b * np, sq, hn] | ||
| 547 | key.transpose(0, 1).transpose(1, 2), # [b * np, hn, sk] | ||
| 548 | beta=0.0, | ||
| 549 | alpha=attn_scale, | ||
| 550 | ) | ||
| 551 | |||
| 552 | # change view to [b, np, sq, sk] | ||
| 553 | attention_scores = matmul_result.view(*output_size) | ||
| 554 | |||
| 555 | if attention_mask is not None: | ||
| 556 | attention_scores.masked_fill_(attention_mask, float('-inf')) | ||
| 557 | # Attention probabilities [b, np, sq, sk] | ||
| 558 | attention_scores = F.softmax( | ||
| 559 | attention_scores, dim=-1, dtype=torch.float32 | ||
| 560 | ) | ||
| 561 | attention_scores = attention_scores.sum(dim=1) | ||
| 562 | if parallel_state.get_tensor_model_parallel_world_size() > 1: | ||
| 563 | # attention scores are scattered to TP ranks in head dimension. | ||
| 564 | torch.distributed.all_reduce(attention_scores.contiguous(), | ||
| 565 | group=parallel_state.get_tensor_model_parallel_group()) | ||
| 566 | return attention_scores | ||
| @@ -5,15 +5,17 @@ from typing import Union | |||
| 5 | 5 | ||
| 6 | import torch | 6 | import torch |
| 7 | import torch.nn.functional as F | 7 | import torch.nn.functional as F |
| 8 | |||
| 8 | from mindspeed.core.context_parallel.ulysses_context_parallel.ulysses_context_parallel import UlyssesContextAttention | 9 | from mindspeed.core.context_parallel.ulysses_context_parallel.ulysses_context_parallel import UlyssesContextAttention |
| 9 | from mindspeed.core.parallel_state import get_context_parallel_group_for_hybrid_ulysses | 10 | from mindspeed.core.parallel_state import get_context_parallel_group_for_hybrid_ulysses |
| 10 | from mindspeed.core.tensor_parallel.random import CheckpointWithoutOutput | 11 | from mindspeed.core.tensor_parallel.random import CheckpointWithoutOutput |
| 11 | from megatron.core.transformer.identity_op import IdentityOp | 12 | from mindspeed.core.transformer.transformer_block import _get_layer_offset |
| 12 | from mindspeed.utils import set_position_ids, get_position_ids | 13 | from mindspeed.utils import set_position_ids, get_position_ids |
| 13 | from mindspeed.core.context_parallel.get_batch_utils import get_actual_seq_len, set_actual_seq_len | 14 | from mindspeed.core.context_parallel.get_batch_utils import get_actual_seq_len, set_actual_seq_len |
| 14 | from mindspeed.core.transformer.moe.moe_feature.fb_overlap.modules.attention import launch_async_all2all_hook, launch_async_all2all | 15 | from mindspeed.core.transformer.moe.moe_feature.fb_overlap.modules.attention import launch_async_all2all_hook, launch_async_all2all |
| 15 | from mindspeed.core.transformer.moe.moe_feature.fb_overlap.modules.utils import TensorSwapManager | 16 | from mindspeed.core.transformer.moe.moe_feature.fb_overlap.modules.utils import TensorSwapManager |
| 16 | 17 | ||
| 18 | from megatron.core.transformer.identity_op import IdentityOp | ||
| 17 | from megatron.core.models.common.embeddings.rotary_pos_embedding import apply_rotary_pos_emb | 19 | from megatron.core.models.common.embeddings.rotary_pos_embedding import apply_rotary_pos_emb |
| 18 | from megatron.core.tensor_parallel import ColumnParallelLinear, RowParallelLinear | 20 | from megatron.core.tensor_parallel import ColumnParallelLinear, RowParallelLinear |
| 19 | from megatron.core.tensor_parallel.mappings import gather_from_sequence_parallel_region | 21 | from megatron.core.tensor_parallel.mappings import gather_from_sequence_parallel_region |
| @@ -25,7 +27,8 @@ from megatron.training import get_args | |||
| 25 | 27 | ||
| 26 | from mindspeed_llm.core.tensor_parallel.layers import LinearNoTP | 28 | from mindspeed_llm.core.tensor_parallel.layers import LinearNoTP |
| 27 | from mindspeed_llm.core.transformer.custom_layers.transformer_engine import PTNorm | 29 | from mindspeed_llm.core.transformer.custom_layers.transformer_engine import PTNorm |
| 28 | from mindspeed_llm.tasks.models.transformer.dsa_indexer import get_dsa_indexer_spec | 30 | from mindspeed_llm.tasks.models.transformer.dsa_indexer import get_dsa_indexer_spec, DSAIndexerLossAutoScaler, \ |
| 31 | compute_dsa_indexer_loss, get_attn_scores, DSAIndexerLossLoggingHelper | ||
| 29 | from mindspeed_llm.tasks.models.transformer.mla_dot_product_attention import MlaDotProductAttention | 32 | from mindspeed_llm.tasks.models.transformer.mla_dot_product_attention import MlaDotProductAttention |
| 30 | from mindspeed_llm.tasks.models.transformer.mla_up_proj_overlap_tp_comm import mla_up_projection_overlap_tp_comm | 33 | from mindspeed_llm.tasks.models.transformer.mla_up_proj_overlap_tp_comm import mla_up_projection_overlap_tp_comm |
| 31 | 34 | ||
| @@ -289,7 +292,7 @@ class CustomMLASelfAttention(SelfAttention): | |||
| 289 | def mla_attention(hidden_states): | 292 | def mla_attention(hidden_states): |
| 290 | args = get_args() | 293 | args = get_args() |
| 291 | tp_size = parallel_state.get_tensor_model_parallel_world_size() | 294 | tp_size = parallel_state.get_tensor_model_parallel_world_size() |
| 292 | 295 | ||
| 293 | # For self attention we just duplicate the rotary_pos_emb if it isn't already | 296 | # For self attention we just duplicate the rotary_pos_emb if it isn't already |
| 294 | nonlocal rotary_pos_emb | 297 | nonlocal rotary_pos_emb |
| 295 | if rotary_pos_emb is not None and not isinstance(rotary_pos_emb, tuple): | 298 | if rotary_pos_emb is not None and not isinstance(rotary_pos_emb, tuple): |
| @@ -446,6 +449,28 @@ class CustomMLASelfAttention(SelfAttention): | |||
| 446 | attention_bias=None, | 449 | attention_bias=None, |
| 447 | packed_seq_params=packed_seq_params, | 450 | packed_seq_params=packed_seq_params, |
| 448 | ) | 451 | ) |
| 452 | if args.enable_dsa_indexer and self.training and torch.is_grad_enabled(): | ||
| 453 | main_attn_dist = get_attn_scores(query, | ||
| 454 | key, | ||
| 455 | attention_mask, | ||
| 456 | self.num_attention_heads_per_partition // | ||
| 457 | self.num_query_groups_per_partition, | ||
| 458 | self.core_attention.scale, | ||
| 459 | ) | ||
| 460 | loss = compute_dsa_indexer_loss( | ||
| 461 | main_attn_dist.detach(), | ||
| 462 | topk_score, | ||
| 463 | topk_indices, | ||
| 464 | args.indexer_loss_coeff, | ||
| 465 | ) | ||
| 466 | |||
| 467 | DSAIndexerLossLoggingHelper.save_loss_to_tracker( | ||
| 468 | loss, | ||
| 469 | _get_layer_offset(args) + self.layer_number, | ||
| 470 | self.config.num_layers, | ||
| 471 | avg_group=parallel_state.get_tensor_and_context_parallel_group(), | ||
| 472 | ) | ||
| 473 | core_attn_out = DSAIndexerLossAutoScaler.apply(core_attn_out, loss) | ||
| 449 | 474 | ||
| 450 | if self.recompute_mla_up_proj_ckpt and core_attn_out.requires_grad: | 475 | if self.recompute_mla_up_proj_ckpt and core_attn_out.requires_grad: |
| 451 | self.recompute_mla_up_proj_ckpt.discard_output() | 476 | self.recompute_mla_up_proj_ckpt.discard_output() |
| @@ -519,12 +544,12 @@ def recompute_mla(mla_checkpoint_manager): | |||
| 519 | change_seq_len = True | 544 | change_seq_len = True |
| 520 | old_actual_seq_len = get_actual_seq_len() | 545 | old_actual_seq_len = get_actual_seq_len() |
| 521 | set_actual_seq_len(actual_seq_len) | 546 | set_actual_seq_len(actual_seq_len) |
| 522 | 547 | ||
| 523 | mla_checkpoint_manager.recompute(grad) | 548 | mla_checkpoint_manager.recompute(grad) |
| 524 | 549 | ||
| 525 | if change_pos_id: | 550 | if change_pos_id: |
| 526 | set_position_ids(old_position_id) | 551 | set_position_ids(old_position_id) |
| 527 | if change_seq_len: | 552 | if change_seq_len: |
| 528 | set_actual_seq_len(old_actual_seq_len) | 553 | set_actual_seq_len(old_actual_seq_len) |
| 529 | 554 | ||
| 530 | return hook_fn | 555 | return hook_fn |
| @@ -28,6 +28,7 @@ import torch | |||
| 28 | import torch_npu | 28 | import torch_npu |
| 29 | 29 | ||
| 30 | from megatron.core.transformer.moe.moe_utils import track_moe_metrics | 30 | from megatron.core.transformer.moe.moe_utils import track_moe_metrics |
| 31 | from megatron.core.transformer.multi_token_prediction import MTPLossLoggingHelper | ||
| 31 | from megatron.training import get_args | 32 | from megatron.training import get_args |
| 32 | from megatron.training import get_timers | 33 | from megatron.training import get_timers |
| 33 | from megatron.training import get_signal_handler | 34 | from megatron.training import get_signal_handler |
| @@ -42,7 +43,8 @@ from megatron.training.checkpointing import save_checkpoint | |||
| 42 | from megatron.training.initialize import initialize_megatron | 43 | from megatron.training.initialize import initialize_megatron |
| 43 | from megatron.training.initialize import write_args_to_tensorboard | 44 | from megatron.training.initialize import write_args_to_tensorboard |
| 44 | from megatron.training.arguments import core_transformer_config_from_args | 45 | from megatron.training.arguments import core_transformer_config_from_args |
| 45 | from megatron.training.training import disable_forward_pre_hook, enable_forward_pre_hook, training_log | 46 | from megatron.training.theoretical_memory_usage import report_theoretical_memory |
| 47 | from megatron.training.training import disable_forward_pre_hook, enable_forward_pre_hook | ||
| 46 | from megatron.training.training import ( | 48 | from megatron.training.training import ( |
| 47 | train_step, calc_params_l2_norm, | 49 | train_step, calc_params_l2_norm, |
| 48 | evaluate_and_print_results, | 50 | evaluate_and_print_results, |
| @@ -53,12 +55,15 @@ from megatron.training.training import ( | |||
| 53 | import megatron.training.utils | 55 | import megatron.training.utils |
| 54 | from megatron.training.utils import ( | 56 | from megatron.training.utils import ( |
| 55 | check_adlr_autoresume_termination, | 57 | check_adlr_autoresume_termination, |
| 58 | reduce_max_stat_across_model_parallel_group, | ||
| 59 | is_last_rank, | ||
| 56 | print_rank_0, | 60 | print_rank_0, |
| 57 | print_rank_last, | 61 | print_rank_last, |
| 58 | report_memory, | 62 | report_memory, |
| 59 | ) | 63 | ) |
| 60 | from megatron.core.distributed import DistributedDataParallel as DDP | 64 | from megatron.core.distributed import DistributedDataParallel as DDP |
| 61 | from megatron.core.distributed import finalize_model_grads | 65 | from megatron.core.distributed import finalize_model_grads |
| 66 | from mindspeed_llm.tasks.models.transformer.dsa_indexer import DSAIndexerLossLoggingHelper | ||
| 62 | from mindspeed_llm.training.initialize import set_jit_fusion_options | 67 | from mindspeed_llm.training.initialize import set_jit_fusion_options |
| 63 | from mindspeed_llm.tasks.posttrain.lora.utils import is_enable_lora | 68 | from mindspeed_llm.tasks.posttrain.lora.utils import is_enable_lora |
| 64 | 69 | ||
| @@ -840,4 +845,278 @@ def num_floating_point_operations_wrapper(fn): | |||
| 840 | if common.zit_scale_in_running_state(): | 845 | if common.zit_scale_in_running_state(): |
| 841 | batch_size = get_args().global_batch_size | 846 | batch_size = get_args().global_batch_size |
| 842 | return fn(args, batch_size) | 847 | return fn(args, batch_size) |
| 843 | return wrapper | 848 | return wrapper |
| 849 | |||
| 850 | |||
| 851 | def training_log(loss_dict, total_loss_dict, learning_rate, decoupled_learning_rate, iteration, | ||
| 852 | loss_scale, report_memory_flag, skipped_iter, | ||
| 853 | grad_norm, params_norm, num_zeros_in_grad): | ||
| 854 | """Log training information such as losses, timing, ....""" | ||
| 855 | args = get_args() | ||
| 856 | timers = get_timers() | ||
| 857 | writer = get_tensorboard_writer() | ||
| 858 | wandb_writer = get_wandb_writer() | ||
| 859 | one_logger = get_one_logger() | ||
| 860 | |||
| 861 | # Advanced, skipped, and Nan iterations. | ||
| 862 | advanced_iters_key = 'advanced iterations' | ||
| 863 | skipped_iters_key = 'skipped iterations' | ||
| 864 | nan_iters_key = 'nan iterations' | ||
| 865 | # Advanced iterations. | ||
| 866 | if not skipped_iter: | ||
| 867 | total_loss_dict[advanced_iters_key] = total_loss_dict.get( | ||
| 868 | advanced_iters_key, 0) + 1 | ||
| 869 | else: | ||
| 870 | if advanced_iters_key not in total_loss_dict: | ||
| 871 | total_loss_dict[advanced_iters_key] = 0 | ||
| 872 | # Skipped iterations. | ||
| 873 | total_loss_dict[skipped_iters_key] = total_loss_dict.get( | ||
| 874 | skipped_iters_key, 0) + skipped_iter | ||
| 875 | # Update losses and set nan iterations | ||
| 876 | got_nan = False | ||
| 877 | for key in loss_dict: | ||
| 878 | if not skipped_iter: | ||
| 879 | total_loss_dict[key] = total_loss_dict.get( | ||
| 880 | key, torch.tensor([0.0], dtype=torch.float, device='cuda')) + loss_dict[key] | ||
| 881 | else: | ||
| 882 | value = loss_dict[key].float().sum().item() | ||
| 883 | is_nan = value == float('inf') or \ | ||
| 884 | value == -float('inf') or \ | ||
| 885 | value != value | ||
| 886 | got_nan = got_nan or is_nan | ||
| 887 | total_loss_dict[nan_iters_key] = total_loss_dict.get( | ||
| 888 | nan_iters_key, 0) + int(got_nan) | ||
| 889 | |||
| 890 | # Logging. | ||
| 891 | timers_to_log = [ | ||
| 892 | 'forward-backward', | ||
| 893 | 'forward-compute', | ||
| 894 | 'backward-compute', | ||
| 895 | 'batch-generator', | ||
| 896 | 'forward-recv', | ||
| 897 | 'forward-send', | ||
| 898 | 'backward-recv', | ||
| 899 | 'backward-send', | ||
| 900 | 'forward-send-forward-recv', | ||
| 901 | 'forward-send-backward-recv', | ||
| 902 | 'backward-send-forward-recv', | ||
| 903 | 'backward-send-backward-recv', | ||
| 904 | 'forward-backward-send-forward-backward-recv', | ||
| 905 | 'layernorm-grads-all-reduce', | ||
| 906 | 'embedding-grads-all-reduce', | ||
| 907 | 'all-grads-sync', | ||
| 908 | 'params-all-gather', | ||
| 909 | 'optimizer-copy-to-main-grad', | ||
| 910 | 'optimizer-unscale-and-check-inf', | ||
| 911 | 'optimizer-clip-main-grad', | ||
| 912 | 'optimizer-count-zeros', | ||
| 913 | 'optimizer-inner-step', | ||
| 914 | 'optimizer-copy-main-to-model-params', | ||
| 915 | 'optimizer'] | ||
| 916 | |||
| 917 | # Calculate batch size. | ||
| 918 | batch_size = args.micro_batch_size * args.data_parallel_size * \ | ||
| 919 | get_num_microbatches() | ||
| 920 | |||
| 921 | # Track app tag & app tag ID | ||
| 922 | one_logger_utils.track_app_tag(batch_size, args.world_size, args.seq_length) | ||
| 923 | |||
| 924 | total_iterations = total_loss_dict[advanced_iters_key] + \ | ||
| 925 | total_loss_dict[skipped_iters_key] | ||
| 926 | |||
| 927 | # learning rate will be None on ranks without trainable params, so we must gather across mp ranks | ||
| 928 | learning_rate = reduce_max_stat_across_model_parallel_group(learning_rate) | ||
| 929 | # Tensorboard values. | ||
| 930 | # Timer requires all the ranks to call. | ||
| 931 | if args.log_timers_to_tensorboard and \ | ||
| 932 | (iteration % args.tensorboard_log_interval == 0): | ||
| 933 | timers.write(timers_to_log, writer, iteration, | ||
| 934 | normalizer=total_iterations) | ||
| 935 | if writer and (iteration % args.tensorboard_log_interval == 0): | ||
| 936 | if wandb_writer: | ||
| 937 | wandb_writer.log({'samples vs steps': args.consumed_train_samples}, | ||
| 938 | iteration) | ||
| 939 | writer.add_scalar('learning-rate', learning_rate, iteration) | ||
| 940 | writer.add_scalar('learning-rate vs samples', learning_rate, | ||
| 941 | args.consumed_train_samples) | ||
| 942 | if wandb_writer: | ||
| 943 | wandb_writer.log({'learning-rate': learning_rate}, iteration) | ||
| 944 | if args.decoupled_lr is not None: | ||
| 945 | writer.add_scalar('decoupled-learning-rate', decoupled_learning_rate, iteration) | ||
| 946 | if args.skipped_train_samples > 0: | ||
| 947 | writer.add_scalar('skipped-train-samples', args.skipped_train_samples, iteration) | ||
| 948 | if wandb_writer: | ||
| 949 | wandb_writer.log({'skipped-train-samples': args.skipped_train_samples}, iteration) | ||
| 950 | writer.add_scalar('batch-size', batch_size, iteration) | ||
| 951 | writer.add_scalar('batch-size vs samples', batch_size, | ||
| 952 | args.consumed_train_samples) | ||
| 953 | if wandb_writer: | ||
| 954 | wandb_writer.log({'batch-size': batch_size}, iteration) | ||
| 955 | for key in loss_dict: | ||
| 956 | writer.add_scalar(key , loss_dict[key], iteration) | ||
| 957 | writer.add_scalar(key + ' vs samples', loss_dict[key], | ||
| 958 | args.consumed_train_samples) | ||
| 959 | if wandb_writer: | ||
| 960 | wandb_writer.log({key: loss_dict[key]}, iteration) | ||
| 961 | if args.log_loss_scale_to_tensorboard: | ||
| 962 | writer.add_scalar('loss-scale', loss_scale, iteration) | ||
| 963 | writer.add_scalar('loss-scale vs samples', loss_scale, | ||
| 964 | args.consumed_train_samples) | ||
| 965 | if wandb_writer: | ||
| 966 | wandb_writer.log({'loss-scale': loss_scale}, iteration) | ||
| 967 | if args.log_world_size_to_tensorboard: | ||
| 968 | writer.add_scalar('world-size', args.world_size, iteration) | ||
| 969 | writer.add_scalar('world-size vs samples', args.world_size, | ||
| 970 | args.consumed_train_samples) | ||
| 971 | if wandb_writer: | ||
| 972 | wandb_writer.log({'world-size': args.world_size}, iteration) | ||
| 973 | if grad_norm is not None: | ||
| 974 | writer.add_scalar('grad-norm', grad_norm, iteration) | ||
| 975 | writer.add_scalar('grad-norm vs samples', grad_norm, | ||
| 976 | args.consumed_train_samples) | ||
| 977 | if wandb_writer: | ||
| 978 | wandb_writer.log({'grad-norm': grad_norm}, iteration) | ||
| 979 | if num_zeros_in_grad is not None: | ||
| 980 | writer.add_scalar('num-zeros', num_zeros_in_grad, iteration) | ||
| 981 | writer.add_scalar('num-zeros vs samples', num_zeros_in_grad, | ||
| 982 | args.consumed_train_samples) | ||
| 983 | if wandb_writer: | ||
| 984 | wandb_writer.log({'num-zeros': num_zeros_in_grad}, iteration) | ||
| 985 | if params_norm is not None: | ||
| 986 | writer.add_scalar('params-norm', params_norm, iteration) | ||
| 987 | writer.add_scalar('params-norm vs samples', params_norm, | ||
| 988 | args.consumed_train_samples) | ||
| 989 | if wandb_writer: | ||
| 990 | wandb_writer.log({'params-norm': params_norm}, iteration) | ||
| 991 | if args.log_memory_to_tensorboard: | ||
| 992 | mem_stats = torch.cuda.memory_stats() | ||
| 993 | writer.add_scalar( | ||
| 994 | "mem-reserved-bytes", | ||
| 995 | mem_stats["reserved_bytes.all.current"], | ||
| 996 | iteration, | ||
| 997 | ) | ||
| 998 | writer.add_scalar( | ||
| 999 | "mem-allocated-bytes", | ||
| 1000 | mem_stats["allocated_bytes.all.current"], | ||
| 1001 | iteration, | ||
| 1002 | ) | ||
| 1003 | writer.add_scalar( | ||
| 1004 | "mem-max-allocated-bytes", | ||
| 1005 | mem_stats["allocated_bytes.all.peak"], | ||
| 1006 | iteration, | ||
| 1007 | ) | ||
| 1008 | writer.add_scalar( | ||
| 1009 | "mem-allocated-count", | ||
| 1010 | mem_stats["allocation.all.current"], | ||
| 1011 | iteration, | ||
| 1012 | ) | ||
| 1013 | if args.num_experts is not None: | ||
| 1014 | moe_loss_scale = 1 / get_num_microbatches() | ||
| 1015 | track_names = [] | ||
| 1016 | if args.moe_router_load_balancing_type in ["aux_loss", "seq_aux_loss"]: | ||
| 1017 | track_names.append("load_balancing_loss") | ||
| 1018 | if args.moe_z_loss_coeff is not None: | ||
| 1019 | track_names.append("z_loss") | ||
| 1020 | track_moe_metrics( | ||
| 1021 | loss_scale=moe_loss_scale, | ||
| 1022 | iteration=iteration, | ||
| 1023 | writer=writer, | ||
| 1024 | wandb_writer=wandb_writer, | ||
| 1025 | total_loss_dict=total_loss_dict, | ||
| 1026 | per_layer_logging=args.moe_per_layer_logging, | ||
| 1027 | force_initialize=True, | ||
| 1028 | track_names=track_names, | ||
| 1029 | num_layers=args.num_layers, | ||
| 1030 | moe_layer_freq=args.moe_layer_freq | ||
| 1031 | ) | ||
| 1032 | if args.mtp_num_layers is not None: | ||
| 1033 | mtp_loss_scale = 1 / get_num_microbatches() | ||
| 1034 | MTPLossLoggingHelper.track_mtp_metrics( | ||
| 1035 | mtp_loss_scale, iteration, writer, wandb_writer, total_loss_dict | ||
| 1036 | ) | ||
| 1037 | if args.enable_dsa_indexer: | ||
| 1038 | dsa_indexer_loss_scale = 1 / get_num_microbatches() | ||
| 1039 | DSAIndexerLossLoggingHelper.track_das_indexer_metrics( | ||
| 1040 | dsa_indexer_loss_scale, iteration, writer, wandb_writer, total_loss_dict | ||
| 1041 | ) | ||
| 1042 | if iteration % args.log_interval == 0: | ||
| 1043 | if args.record_memory_history and is_last_rank(): | ||
| 1044 | snapshot = torch.cuda.memory._snapshot() | ||
| 1045 | from pickle import dump | ||
| 1046 | with open(args.memory_snapshot_path, 'wb') as f: | ||
| 1047 | dump(snapshot, f) | ||
| 1048 | |||
| 1049 | elapsed_time = timers('interval-time').elapsed(barrier=True) | ||
| 1050 | elapsed_time_per_iteration = elapsed_time / total_iterations | ||
| 1051 | |||
| 1052 | throughput = num_floating_point_operations(args, batch_size) / ( | ||
| 1053 | elapsed_time_per_iteration * 10**12 * args.world_size) | ||
| 1054 | |||
| 1055 | one_logger_utils.track_e2e_metrics(args.log_throughput, throughput) | ||
| 1056 | |||
| 1057 | if args.log_timers_to_tensorboard: | ||
| 1058 | if writer: | ||
| 1059 | writer.add_scalar('iteration-time', | ||
| 1060 | elapsed_time_per_iteration, iteration) | ||
| 1061 | if wandb_writer: | ||
| 1062 | wandb_writer.log({'iteration-time': elapsed_time_per_iteration}, | ||
| 1063 | iteration) | ||
| 1064 | log_string = f" [{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}]" | ||
| 1065 | log_string += ' iteration {:8d}/{:8d} |'.format( | ||
| 1066 | iteration, args.train_iters) | ||
| 1067 | log_string += ' consumed samples: {:12d} |'.format( | ||
| 1068 | args.consumed_train_samples) | ||
| 1069 | if args.skipped_train_samples > 0: | ||
| 1070 | log_string += ' skipped samples: {:12d} |'.format( | ||
| 1071 | args.skipped_train_samples) | ||
| 1072 | log_string += ' elapsed time per iteration (ms): {:.1f} |'.format( | ||
| 1073 | elapsed_time_per_iteration * 1000.0) | ||
| 1074 | if args.log_throughput: | ||
| 1075 | log_string += f' throughput per GPU (TFLOP/s/GPU): {throughput:.1f} |' | ||
| 1076 | if args.log_timers_to_tensorboard: | ||
| 1077 | if writer: | ||
| 1078 | writer.add_scalar('throughput', throughput, iteration) | ||
| 1079 | if wandb_writer: | ||
| 1080 | wandb_writer.log({'throughput': throughput}, iteration) | ||
| 1081 | # Decoupled_learning_rate should be not None only on first and last pipeline stage. | ||
| 1082 | log_string += f' learning rate: {learning_rate:.6E} |' | ||
| 1083 | if args.decoupled_lr is not None and (mpu.is_pipeline_first_stage(ignore_virtual=True) or | ||
| 1084 | mpu.is_pipeline_last_stage(ignore_virtual=True)): | ||
| 1085 | assert decoupled_learning_rate is not None | ||
| 1086 | log_string += f' decoupled learning rate: {decoupled_learning_rate:.6E} |' | ||
| 1087 | else: | ||
| 1088 | assert decoupled_learning_rate is None | ||
| 1089 | log_string += f' global batch size: {batch_size:5d} |' | ||
| 1090 | for key in total_loss_dict: | ||
| 1091 | if key not in [advanced_iters_key, skipped_iters_key, | ||
| 1092 | nan_iters_key]: | ||
| 1093 | avg = total_loss_dict[key].item() / \ | ||
| 1094 | float(max(1, total_loss_dict[advanced_iters_key])) | ||
| 1095 | if avg > 0.0: | ||
| 1096 | log_string += ' {}: {:.6E} |'.format(key, avg) | ||
| 1097 | total_loss_dict[key] = torch.tensor([0.0], dtype=torch.float, device='cuda') | ||
| 1098 | log_string += f' loss scale: {loss_scale:.1f} |' | ||
| 1099 | if grad_norm is not None: | ||
| 1100 | log_string += f' grad norm: {grad_norm:.3f} |' | ||
| 1101 | if num_zeros_in_grad is not None: | ||
| 1102 | log_string += f' num zeros: {num_zeros_in_grad} |' | ||
| 1103 | if params_norm is not None: | ||
| 1104 | log_string += f' params norm: {params_norm:.3f} |' | ||
| 1105 | log_string += ' number of skipped iterations: {:3d} |'.format( | ||
| 1106 | total_loss_dict[skipped_iters_key]) | ||
| 1107 | log_string += ' number of nan iterations: {:3d} |'.format( | ||
| 1108 | total_loss_dict[nan_iters_key]) | ||
| 1109 | total_loss_dict[advanced_iters_key] = 0 | ||
| 1110 | total_loss_dict[skipped_iters_key] = 0 | ||
| 1111 | total_loss_dict[nan_iters_key] = 0 | ||
| 1112 | print_rank_last(log_string) | ||
| 1113 | if report_memory_flag: | ||
| 1114 | # Report memory after optimizer state has been initialized. | ||
| 1115 | if torch.distributed.get_rank() == 0: | ||
| 1116 | num_microbatches = get_num_microbatches() | ||
| 1117 | report_theoretical_memory(args, num_microbatches=num_microbatches, verbose=True) | ||
| 1118 | report_memory(f'(after {iteration} iterations)') | ||
| 1119 | report_memory_flag = False | ||
| 1120 | timers.log(timers_to_log, normalizer=args.log_interval) | ||
| 1121 | |||
| 1122 | return report_memory_flag | ||
| @@ -6,6 +6,11 @@ from tests.mindspore.test_tools.acquire_json import transfer_logs_as_json, read_ | |||
| 6 | LOSS = "lm loss" | 6 | LOSS = "lm loss" |
| 7 | 7 | ||
| 8 | 8 | ||
| 9 | class TestMargin: | ||
| 10 | _MARGIN_NAME = " margin" | ||
| 11 | loss = 0.02 | ||
| 12 | |||
| 13 | |||
| 9 | class TestCIST: | 14 | class TestCIST: |
| 10 | 15 | ||
| 11 | 16 | ||
| @@ -56,7 +61,8 @@ class TestCIST: | |||
| 56 | def _compare_lm_loss(self, expected_list, actual_list): | 61 | def _compare_lm_loss(self, expected_list, actual_list): |
| 57 | for step, (expected_val, actual_val) in enumerate(zip(expected_list, actual_list)): | 62 | for step, (expected_val, actual_val) in enumerate(zip(expected_list, actual_list)): |
| 58 | print(f"Checking step {step + 1} for lm loss") | 63 | print(f"Checking step {step + 1} for lm loss") |
| 59 | assert actual_val == expected_val, f"The loss at step {step} should be {expected_val} but it is {actual_val}." | 64 | assert actual_val == pytest.approx(expected=expected_val, rel=TestMargin.loss), \ |
| 65 | f"The loss at step {step} should be approximate to {expected_val} but it is {actual_val}." | ||
| 60 | 66 | ||
| 61 | def test_lm_loss(self, baseline_json, generate_log, generate_json): | 67 | def test_lm_loss(self, baseline_json, generate_log, generate_json): |
| 62 | # expected training loss curve at different global steps. | 68 | # expected training loss curve at different global steps. |