Commit bf97f933 authored Jul 24, 2020 by Jared Casper

Merge branch 'apex-fix' into 'master'

Initialization fixes: allowing simple case like pytest pass, also making apex optional

See merge request ADLR/megatron-lm!99

parents 05620ee4 b04eb0a2

megatron/initialize.py

+6 −0

Original line number	Diff line number	Diff line
		@@ -39,6 +39,12 @@ def initialize_megatron(extra_args_provider=None, args_defaults={},
		# Make sure cuda is available.
		assert torch.cuda.is_available(), 'Megatron requires CUDA.'

		# This is temporary WAR to make simple case like pytest calling with same args twice
		# Need to implement clean factory init.
		if mpu.model_parallel_is_initialized():
		return


		# Parse args, build tokenizer, and set adlr-autoresume,
		# tensorboard-writer, and timers.
		set_global_variables(extra_args_provider=extra_args_provider,

megatron/model/transformer.py

+8 −1

Original line number	Diff line number	Diff line
		@@ -18,7 +18,14 @@
		import math

		import torch
		try:
		from apex.normalization.fused_layer_norm import FusedLayerNorm as LayerNorm
		# Try to use FusedLayerNorm from Apex - this will trigger an error.
		_ = FusedLayerNorm(8, eps=1e-5)

		except Exception as e:
		print('WARNING: APEX is not available, using torch.nn.LayerNorm instead of apex.normalization.FusedLayerNorm!')
		from torch.nn import LayerNorm

		from megatron import get_args
		from megatron import mpu