[embed norm] switch to apex MixedFusedLayerNorm

[stas00]

as noticed by @thomasw21 - switching embed layernorm to use MixedFusedLayerNorm for consistency with other layer norms.

---

Incidentally, this also fixes a bug with how torch.nn.LayerNorm was used until now.

the framework was putting LayerNorm into the wrong param group here:

Megatron-DeepSpeed/megatron/optimizer/__init__.py

Lines 31 to 45 in dd06ea3

	weight_decay_params = {'params': []}
	no_weight_decay_params = {'params': [], 'weight_decay': 0.0}
	for module in modules:
	for module_ in module.modules():
	if isinstance(module_, LayerNorm):
	no_weight_decay_params['params'].extend(
	[p for p in list(module_._parameters.values())
	if p is not None])
	else:
	weight_decay_params['params'].extend(
	[p for n, p in list(module_._parameters.items())
	if p is not None and n != 'bias'])
	no_weight_decay_params['params'].extend(
	[p for n, p in list(module_._parameters.items())
	if p is not None and n == 'bias'])

it should have been in no_weight_decay_params but ended up in weight_decay_params because in this module LayerNorm is an alias for MixedFusedLayerNorm, so if isinstance(module_, LayerNorm) was False.

So if we want to use torch.nn.LayerNorm we have to change the code above to additionally check for or isinstance(module_, torch.nn.LayerNorm).