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

        # 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,

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:

    from torch.nn import LayerNorm

from megatron import get_args

from megatron import mpu