move utils to picotron

This commit is contained in:
ferdinand.mom 2024-11-04 15:32:20 +00:00
parent 8af19d0caa
commit a5706858e0
2 changed files with 1 additions and 1 deletions

View File

@ -22,7 +22,7 @@ from transformers import AutoConfig
import numpy as np
from picotron.parallel.tensor_parallel.tensor_parallel import TensorParallel
import picotron.process_group_manager as pgm
from utils import MicroBatchDataLoader, set_all_seed, print, to_readable_format, save_checkpoint, load_checkpoint
from picotron.parallel.utils import MicroBatchDataLoader, set_all_seed, print, to_readable_format, save_checkpoint, load_checkpoint
from picotron.process_group_manager import setup_process_group_manager
from picotron.parallel.pipeline_parallel import train_step_pipeline_1f1b, train_step_pipeline_afab, PipelineParallel
from picotron.parallel.data_parallel.data_parallel_bucket import DataParallel