diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py
index 652333ac..6f4330a3 100644
--- a/megatron/training/checkpointing.py
+++ b/megatron/training/checkpointing.py
@@ -255,12 +255,14 @@ def read_metadata(tracker_filename):
                 print_rank_0('ERROR: Invalid metadata file {}. Exiting'.format(
                     tracker_filename))
                 sys.exit()
-    assert iteration > 0 or release, 'error parsing metadata file {}'.format(
+    # TODO: we use iteration 0 to load checkpoint from other framework.  
+    # We should remove this after we have a better way to load checkpoint from other framework.
+    assert iteration >= 0 or release, 'error parsing metadata file {}'.format(
         tracker_filename)
 
     # Get the max iteration retrieved across the ranks.
     if torch.distributed.is_initialized():
-        iters_cuda = torch.tensor([iteration], dtype=torch.long, device='cuda')
+        iters_cuda = torch.tensor([iteration], dtype=torch.long, device='cuda' if 'nccl' in torch.distributed.get_backend() else 'cpu')
         torch.distributed.all_reduce(iters_cuda, op=torch.distributed.ReduceOp.MAX)
         max_iter = iters_cuda[0].item()
 
@@ -592,6 +594,28 @@ def save_checkpoint(iteration, model, optimizer, opt_param_scheduler, num_floati
     if not torch.distributed.is_initialized() \
        or is_last_rank():
         def wandb_finalize_fn():
+            ######### FlagScale Begin #########
+            #NOTE(lizhiyu): The tracker file is created by rank 0 but wandb_finalize_fn is called on the last rank.
+            import time as pytime
+
+            tracker_file = get_checkpoint_tracker_filename(save_dir)
+
+            timeout_seconds = 600  # 10 minutes
+            wait_interval_seconds = 5
+            max_retries = timeout_seconds // wait_interval_seconds
+
+            for _ in range(max_retries):
+                if isfile(tracker_file):
+                    with open(tracker_file, 'r') as f:
+                        content = f.read().strip()
+                        if content == str(iteration):
+                            break  # Success
+                print(f'WandB finalization waiting for the tracker file {tracker_file} to update...')
+                pytime.sleep(wait_interval_seconds)
+            else:
+                # This block executes if the loop completes without a `break`.
+                raise RuntimeError(f"Timed out waiting for tracker file {tracker_file} to be updated for iteration {iteration} after {timeout_seconds} seconds.")
+            ######### FlagScale End #########
             wandb_utils.on_save_checkpoint_success(checkpoint_name, get_checkpoint_tracker_filename(save_dir), save_dir, iteration)
         if args.async_save:
             assert async_save_request is not None
@@ -674,9 +698,7 @@ def maybe_save_dataloader_state(train_iterator, iteration, dataloader_save_path)
 
     torch.distributed.barrier(group=mpu.get_data_parallel_group())
 
-    if mpu.get_data_parallel_rank() == 0:
-        ensure_directory_exists(data_state_save_path)
-
+    ensure_directory_exists(data_state_save_path)
     torch.distributed.barrier(group=mpu.get_data_parallel_group())
 
     dataloader_save_dict = {}
@@ -1095,6 +1117,10 @@ def load_args_from_checkpoint(
             checkpoint_args, 'add_bias_linear', not getattr(checkpoint_args, 'disable_bias_linear')
         )
 
+    # For backward compatibility.
+    if hasattr(checkpoint_args, 'apply_layernorm_rms'):
+        checkpoint_args.normalization = 'RMSNorm'
+
     def _set_arg(arg_name, old_arg_name=None, force=False):
         if not force and getattr(args, arg_name, None) is not None:
             return
@@ -1130,6 +1156,8 @@ def load_args_from_checkpoint(
     _set_arg('add_qkv_bias', force=True)
     _set_arg('squared_relu', force=True)
     _set_arg('swiglu', force=True)
+    _set_arg('multiple_of', force=True)
+    _set_arg('hidden_dim_multiplier', force=True)
     _set_arg('untie_embeddings_and_output_weights', force=True)
     _set_arg('apply_layernorm_1p', force=True)
     _set_arg('normalization', force=True)
