More ict_merge changes and interactive testing

tasks/race/finetune.py

@@ -15,9 +15,8 @@
 
 """Race."""
 
-from megatron import get_args
+from megatron import get_args, print_rank_0
 from megatron import get_tokenizer
-from megatron import print_rank_0
 from megatron.model.multiple_choice import MultipleChoice
 from tasks.eval_utils import accuracy_func_provider
 from tasks.finetune_utils import finetune

tasks/zeroshot_gpt2/datasets.py

@@ -21,9 +21,8 @@ import math
 import numpy as np
 import torch
 
-from megatron import get_args
+from megatron import get_args, print_rank_0
 from megatron import get_tokenizer
-from megatron import print_rank_0
 from .detokenizer import get_detokenizer
 
 

tasks/zeroshot_gpt2/evaluate.py

@@ -19,10 +19,9 @@ import math
 
 import torch
 
-from megatron import get_args
+from megatron import get_args, print_rank_0
 from megatron import get_tokenizer
 from megatron import mpu
-from megatron import print_rank_0
 from megatron.checkpointing import load_checkpoint
 from megatron.model import GPT2Model
 from megatron.training import get_model

tools/generate_samples_gpt2.py

@@ -20,9 +20,8 @@ import sys
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__),
                                              os.path.pardir)))
 
-from megatron import get_args
+from megatron import get_args, print_rank_0
 from megatron import get_tokenizer
-from megatron import print_rank_0
 from megatron.checkpointing import load_checkpoint
 from megatron.initialize import initialize_megatron
 from megatron.model import GPT2Model