Commit 3262398d authored by Raul Puri's avatar Raul Puri
Browse files

Merge branch 'args_fix' into 'master'

args bug fix

See merge request ADLR/megatron-lm!77
parents 3af90c7a 04aa357e
...@@ -47,6 +47,19 @@ def parse_args(extra_args_provider=None, defaults={}, ...@@ -47,6 +47,19 @@ def parse_args(extra_args_provider=None, defaults={},
else: else:
args = parser.parse_args() args = parser.parse_args()
# Distributed args.
args.rank = int(os.getenv('RANK', '0'))
args.world_size = int(os.getenv("WORLD_SIZE", '1'))
args.model_parallel_size = min(args.model_parallel_size, args.world_size)
if args.rank == 0:
print('using world size: {} and model-parallel size: {} '.format(
args.world_size, args.model_parallel_size))
# Fp16 loss scaling.
args.dynamic_loss_scale = False
if args.loss_scale is None:
args.dynamic_loss_scale = True
# Set input defaults. # Set input defaults.
for key in defaults: for key in defaults:
# For default to be valid, it should not be provided in the # For default to be valid, it should not be provided in the
...@@ -67,19 +80,6 @@ def parse_args(extra_args_provider=None, defaults={}, ...@@ -67,19 +80,6 @@ def parse_args(extra_args_provider=None, defaults={},
for req_arg in required_args: for req_arg in required_args:
_check_arg_is_not_none(args, req_arg) _check_arg_is_not_none(args, req_arg)
# Distributed args.
args.rank = int(os.getenv('RANK', '0'))
args.world_size = int(os.getenv("WORLD_SIZE", '1'))
args.model_parallel_size = min(args.model_parallel_size, args.world_size)
if args.rank == 0:
print('using world size: {} and model-parallel size: {} '.format(
args.world_size, args.model_parallel_size))
# Fp16 loss scaling.
args.dynamic_loss_scale = False
if args.loss_scale is None:
args.dynamic_loss_scale = True
# Checks. # Checks.
assert args.hidden_size % args.num_attention_heads == 0 assert args.hidden_size % args.num_attention_heads == 0
if args.seq_length is not None: if args.seq_length is not None:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment