






"""simple docstring"""
import importlib
import sys
from argparse import REMAINDER, ArgumentParser
from pathlib import Path

import torch_xla.distributed.xla_multiprocessing as xmp






def 			lowerCamelCase_    (		):
 lowerCamelCase__	:     Any								=		ArgumentParser(
     description=(
         'PyTorch TPU distributed training launch helper utility that will spawn up multiple distributed processes'
     )							)

 # Optional arguments for the launch helper
 parser.add_argument('--num_cores'     ,  type=_lowerCamelCase     ,  default=1     ,  help='Number of TPU cores to use (1 or 8).'							)

 # positional
 parser.add_argument(
     'training_script'     ,  type=_lowerCamelCase     ,  help=(
         'The full path to the single TPU training '
         'program/script to be launched in parallel, '
         'followed by all the arguments for the '
         'training script'
     )     ,  )

 # rest from the training program
 parser.add_argument('training_script_args'     ,  nargs=_lowerCamelCase							)

 return parser.parse_args()






def 			lowerCamelCase_    (		):
 lowerCamelCase__	:     str								=		parse_args()

 # Import training_script as a module.
 lowerCamelCase__	:     Dict								=		Path(args.training_script							)
 sys.path.append(str(script_fpath.parent.resolve()							)							)
 lowerCamelCase__	:     Optional[Any]								=		script_fpath.stem
 lowerCamelCase__	:     Optional[int]								=		importlib.import_module(_lowerCamelCase							)

 # Patch sys.argv
 lowerCamelCase__	:     List[str]								=		[args.training_script] + args.training_script_args + ['--tpu_num_cores', str(args.num_cores							)]

 xmp.spawn(mod._mp_fn     ,  args=()     ,  nprocs=args.num_cores							)


if __name__ == "__main__":
      main()









"""simple docstring"""
import torch.nn as nn

from .imports import is_fpa_available


if is_fpa_available():
      import transformer_engine.pytorch as te






def 			lowerCamelCase_    (		_lowerCamelCase     ,  _lowerCamelCase=True     ,  _lowerCamelCase=True     ,  _lowerCamelCase=True							):
 if not is_fpa_available():
  raise ImportError('Using `convert_model` requires transformer_engine to be installed.'							)
 for name, module in model.named_children():
  if isinstance(_lowerCamelCase     ,  nn.Linear							) and to_transformer_engine and _convert_linear:
   # Return early if the linear layer weights are not multiples of 16
   if any(p % 16 != 0 for p in module.weight.shape							):
    return
   lowerCamelCase__	:     List[Any]								=		module.bias is not None
   lowerCamelCase__	:     Tuple								=		te.Linear(
       module.in_features     ,  module.out_features     ,  bias=_lowerCamelCase     ,  params_dtype=module.weight.dtype							)
   lowerCamelCase__	:     Optional[Any]								=		module.weight.data.clone()
   if has_bias:
    lowerCamelCase__	:     List[str]								=		module.bias.data.clone()

   setattr(_lowerCamelCase     ,  _lowerCamelCase     ,  _lowerCamelCase							)
  elif isinstance(_lowerCamelCase     ,  nn.LayerNorm							) and to_transformer_engine and _convert_ln:
   lowerCamelCase__	:     Tuple								=		te.LayerNorm(module.normalized_shape[0]     ,  eps=module.eps     ,  params_dtype=module.weight.dtype							)
   lowerCamelCase__	:     Tuple								=		module.weight.data.clone()
   lowerCamelCase__	:     Optional[Any]								=		module.bias.data.clone()

   setattr(_lowerCamelCase     ,  _lowerCamelCase     ,  _lowerCamelCase							)
  elif isinstance(_lowerCamelCase     ,  te.Linear							) and not to_transformer_engine and _convert_linear:
   lowerCamelCase__	:     Optional[Any]								=		module.bias is not None
   lowerCamelCase__	:     Union[str, Any]								=		nn.Linear(
       module.in_features     ,  module.out_features     ,  bias=_lowerCamelCase     ,  params_dtype=module.weight.dtype							)
   lowerCamelCase__	:     str								=		module.weight.data.clone()
   if has_bias:
    lowerCamelCase__	:     Optional[int]								=		module.bias.data.clone()

   setattr(_lowerCamelCase     ,  _lowerCamelCase     ,  _lowerCamelCase							)
  elif isinstance(_lowerCamelCase     ,  te.LayerNorm							) and not to_transformer_engine and _convert_ln:
   lowerCamelCase__	:     str								=		nn.LayerNorm(module.normalized_shape[0]     ,  eps=module.eps     ,  params_dtype=module.weight.dtype							)
   lowerCamelCase__	:     Dict								=		module.weight.data.clone()
   lowerCamelCase__	:     Dict								=		module.bias.data.clone()

   setattr(_lowerCamelCase     ,  _lowerCamelCase     ,  _lowerCamelCase							)
  else:
   convert_model(
       _lowerCamelCase     ,  to_transformer_engine=_lowerCamelCase     ,  _convert_linear=_convert_linear     ,  _convert_ln=_convert_ln     ,  )






def 			lowerCamelCase_    (		_lowerCamelCase							):
 if not is_fpa_available():
  raise ImportError('Using `has_transformer_engine_layers` requires transformer_engine to be installed.'							)
 for m in model.modules():
  if isinstance(_lowerCamelCase     ,  (te.LayerNorm, te.Linear)							):
   return True
 return False

