update dewcriptions being printed

EleutherAI · May 2, 2024 · fabb45b · fabb45b
1 parent 055596f
commit fabb45b
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/megatron/mup_substitute.py b/megatron/mup_substitute.py
@@ -51,12 +51,11 @@ def get_coord_data(
  for i in range(nseeds):
  seed = (i + 1) * 100000
  torch.manual_seed(seed)
- print_rank_0(f">>> muP Coord Check: mup_width_multiplier set to {neox_args.mup_width_multiplier}")
- print_rank_0(f">>> muP Coord Check: Running Model with width: {width} on seed: {seed}\n")
+
  model, optimizer, lr_scheduler = model_obj()
  model.train()
- neox_args.hidden_size = width
-
+ print_rank_0(f">>> muP Coord Check: Running Model with width: {width} on seed: {seed}")
+ print_rank_0(f">>> muP Coord Check: mup_width_multiplier set to {model.neox_args.mup_width_multiplier}")
  for step in range(nsteps + 1):
 
  word_embedding_act_abs_std_list = []