perf-tp-dp-comm-overlap

Par nvidia · skills

Guide opérationnel pour activer le chevauchement des communications TP, DP et PP dans Megatron-Bridge, incluant les paramètres de configuration, les points d'ancrage dans le code, les pièges courants et les procédures de vérification.

npx skills add https://github.com/nvidia/skills --skill perf-tp-dp-comm-overlap

Compétence de Chevauchement de Communication TP / DP / PP

Pour le contexte de base et le niveau de recommandation, voir :

  • @docs/training/communication-overlap.md

Activation

Remplacement minimal de Bridge :

from megatron.bridge.training.comm_overlap import CommOverlapConfig

cfg.model.tensor_model_parallel_size = 4
cfg.model.sequence_parallel = True
cfg.model.pipeline_model_parallel_size = 4
cfg.model.virtual_pipeline_model_parallel_size = 2

cfg.comm_overlap = CommOverlapConfig(
    tp_comm_overlap=True,
)

cfg.ddp.use_distributed_optimizer = True
cfg.ddp.overlap_grad_reduce = True
cfg.ddp.overlap_param_gather = True

Préset TP optionnel :

from megatron.bridge.training.comm_overlap import userbuffers_bf16_h100_h12288_tp4_mbs1_seqlen2048

cfg.comm_overlap.tp_comm_overlap_cfg = userbuffers_bf16_h100_h12288_tp4_mbs1_seqlen2048

Les paramètres de précision appartiennent à la précision mixte :

cfg.mixed_precision.grad_reduce_in_fp32 = False
cfg.mixed_precision.fp8_param_gather = False

Ancres de Code

Portail de chevauchement Bridge :

if self.user_comm_overlap_cfg.tp_comm_overlap is True:
    if model_cfg.tensor_model_parallel_size < 2:
        ...
    elif not model_cfg.sequence_parallel:
        ...
    elif not HAVE_TE:
        ...

Sélection de chevauchement PP :

if model_cfg.pipeline_model_parallel_size > 1:
    if vp_size > 1:
        comm_overlap_cfg.overlap_p2p_comm = True
        comm_overlap_cfg.batch_p2p_comm = False
    else:
        comm_overlap_cfg.overlap_p2p_comm = False
        comm_overlap_cfg.batch_p2p_comm = True

Valeurs par défaut de chevauchement DP :

if self.data_parallel_size > 1:
    comm_overlap_cfg.bucket_size = 128 * 1024 * 1024
    comm_overlap_cfg.overlap_grad_reduce = True
    comm_overlap_cfg.overlap_param_gather = True

Ajustement d'environnement au lancement :

executor.env_vars["CUDA_DEVICE_MAX_CONNECTIONS"] = str(cuda_device_max_connections)
...
executor.env_vars["NVTE_FWD_LAYERNORM_SM_MARGIN"] = str(self.layernorm_sm_margin)
executor.env_vars["NVTE_BWD_LAYERNORM_SM_MARGIN"] = str(self.layernorm_sm_margin)

Pièges

  1. Le chevauchement TP se désactive silencieusement si sequence_parallel=False ou si Transformer Engine n'est pas disponible.
  2. Le chevauchement PP n'est pas activé pour tous les cas PP. Bridge sélectionne automatiquement overlap_p2p_comm=True uniquement quand PP > 1 et VPP > 1.
  3. bucket_size est un paramètre basé sur le nombre de paramètres, pas sur la taille en octets.
  4. grad_reduce_in_fp32 et fp8_param_gather doivent être définis via la précision mixte, et non comme premier réglage DDP autonome.
  5. CUDA_DEVICE_MAX_CONNECTIONS et la marge SM de LayerNorm sont des paramètres de plugin au lancement, pas des champs CommOverlapConfig.

Vérification

Utilisez d'abord la couverture de test unitaire de chevauchement validée :

uv run python -m pytest tests/unit_tests/training/test_comm_overlap.py -q

Vérification optionnelle supplémentaire si nemo_run est disponible :

uv run python -m pytest tests/unit_tests/recipes/test_run_plugins.py -q

Critères de succès :

  • la première commande rapporte 26 passed
  • la deuxième commande valide le câblage env détenu par le plugin quand il n'est pas ignoré

Skills similaires