Skip to content

Commit

Permalink
update mcore changes
Browse files Browse the repository at this point in the history
Signed-off-by: Jimmy Zhang <jiemingz@nvidia.com>
  • Loading branch information
jiemingz committed Aug 28, 2024
1 parent a0c0445 commit 9de6189
Show file tree
Hide file tree
Showing 2 changed files with 3 additions and 8 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@
try:
from megatron.core import parallel_state, tensor_parallel
from megatron.core.fusions.fused_layer_norm import FusedLayerNorm
from megatron.core.transformer.graphs import CudaGraphManager
from megatron.core.transformer.cuda_graphs import CudaGraphManager
from megatron.core.transformer.spec_utils import ModuleSpec
from megatron.core.transformer.transformer_block import TransformerBlockSubmodules, get_num_layers_to_build
from megatron.core.transformer.transformer_layer import BaseTransformerLayer
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -835,14 +835,9 @@ def training_step(self, dataloader_iter):
if not self.mcore_gpt:
module = module.language_model

# Cudagraphed model does not trigger param sync hooks, so manually trigger param syncs here.
if self.cfg.get('enable_cuda_graph', False):
for param in module.parameters():
if hasattr(module, 'embedding'):
for param in module.embedding.parameters():
param.data_ptr()
else:
if hasattr(module, 'embedding'):
for param in module.embedding.parameters():
param.data_ptr()

if self.cfg.get('pipeline_model_parallel_size', 1) > 1 and parallel_state.is_pipeline_last_stage(
ignore_virtual=True
Expand Down

0 comments on commit 9de6189

Please sign in to comment.