diff to current megatron

megatron/__init__.py

@@ -1,26 +1,20 @@
-# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http:https://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-import torch
-
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
 
-def print_rank_0(*message):
- """If distributed is initialized print only on rank 0."""
- if torch.distributed.is_initialized():
- if torch.distributed.get_rank() == 0:
- print(*message, flush=True)
- else:
- print(*message, flush=True)
+import torch
 
+from .global_vars import get_args, get_retro_args
+from .global_vars import get_current_global_batch_size
+from .global_vars import get_num_microbatches
+from .global_vars import get_signal_handler
+from .global_vars import update_num_microbatches
+from .global_vars import get_tokenizer
+from .global_vars import get_tensorboard_writer
+from .global_vars import get_wandb_writer
+from .global_vars import get_one_logger
+from .global_vars import get_adlr_autoresume
+from .global_vars import get_timers
+from .initialize import initialize_megatron
 
-from .neox_arguments import NeoXArgs
+from .utils import (print_rank_0,
+ is_last_rank,
+ print_rank_last)