Skip to content

Commit bcdec34

Browse files
committed
polish
1 parent dc411d9 commit bcdec34

11 files changed

Lines changed: 19 additions & 32 deletions

File tree

flagscale/train/__init__.py

Lines changed: 0 additions & 4 deletions
This file was deleted.

flagscale/train/global_vars.py

Lines changed: 0 additions & 13 deletions
This file was deleted.

flagscale/train/datasets/concated_indexed_dataset.py renamed to flagscale/train/megatron/training/datasets/concated_indexed_dataset.py

File renamed without changes.

flagscale/train/extra_valid.py renamed to flagscale/train/megatron/training/extra_valid.py

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,11 +9,9 @@
99
from megatron.core.rerun_state_machine import RerunDataIterator
1010
from megatron.training.datasets.data_samplers import build_pretraining_data_loader
1111
from megatron.training import get_args, get_tokenizer, print_rank_0
12-
from megatron.training.global_vars import get_tensorboard_writer, get_wandb_writer
12+
from megatron.training.global_vars import get_tensorboard_writer, get_wandb_writer, get_extra_valid_datasets, set_extra_valid_datasets
1313
from megatron.training.utils import is_last_rank, print_rank_last
1414

15-
from flagscale.train import get_extra_valid_datasets, set_extra_valid_datasets
16-
1715

1816
def is_dataset_built_on_rank():
1917
return (

flagscale/train/megatron/training/global_vars.py

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,8 @@
2929
_GLOBAL_SIGNAL_HANDLER = None
3030

3131
_GLOBAL_SPIKY_LOSS_DETECTOR = None
32+
_GLOBAL_EXTRA_VALID_DATASETS = None
33+
3234

3335

3436
def get_args():
@@ -365,7 +367,6 @@ def destroy_global_vars():
365367
_GLOBAL_SIGNAL_HANDLER = None
366368

367369

368-
369370
def get_spiky_loss_detector():
370371
"""Return spiky loss detector."""
371372
_ensure_var_is_initialized(_GLOBAL_SPIKY_LOSS_DETECTOR, "spiky loss detector")
@@ -377,3 +378,14 @@ def set_spiky_loss_detector(args):
377378
global _GLOBAL_SPIKY_LOSS_DETECTOR
378379
_ensure_var_is_not_initialized(_GLOBAL_SPIKY_LOSS_DETECTOR, "spiky loss detector")
379380
_GLOBAL_SPIKY_LOSS_DETECTOR = SpikyLossDetector(args.spiky_loss_threshold)
381+
382+
383+
def get_extra_valid_datasets():
384+
"""Return extra_valid datasets.""" ""
385+
return _GLOBAL_EXTRA_VALID_DATASETS
386+
387+
388+
def set_extra_valid_datasets(extra_valid_datasets):
389+
"""Set extra_valid datasets.""" ""
390+
global _GLOBAL_EXTRA_VALID_DATASETS
391+
_GLOBAL_EXTRA_VALID_DATASETS = extra_valid_datasets

flagscale/train/train.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -131,8 +131,8 @@
131131

132132
from megatron.training import ft_integration
133133

134-
from flagscale.train.extra_valid import extra_evaluate_and_print_results
135-
from flagscale.train.extra_valid import build_extra_valid_data_iterators
134+
from megatron.training.extra_valid import extra_evaluate_and_print_results
135+
from megatron.training.extra_valid import build_extra_valid_data_iterators
136136
from flagscale.train.stablelm2_scheduler import StableLM2SchedulerConfig
137137
from megatron.training.global_vars import get_spiky_loss_detector
138138
from flagscale.train.theoretical_memory_usage import report_theoretical_memory as fs_report_theoretical_memory

flagscale/train/train_aquila_sft.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -45,8 +45,8 @@
4545
except ImportError:
4646
has_nvidia_modelopt = False
4747

48-
from flagscale.train.datasets.sft_dataset import SFTDatasetConfig, SFTDataset
49-
from flagscale.train.extra_valid import extra_valid_datasets_provider
48+
from megatron.training.datasets.sft_dataset_fs import SFTDatasetConfig, SFTDataset
49+
from megatron.training.extra_valid import extra_valid_datasets_provider
5050
from flagscale.train.train import pretrain
5151
from megatron.plugin.hetero.parallel_context import get_parallel_context
5252

flagscale/train/train_gpt.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@
3434
except ImportError:
3535
has_nvidia_modelopt = False
3636

37-
from flagscale.train.extra_valid import extra_valid_datasets_provider
37+
from megatron.training.extra_valid import extra_valid_datasets_provider
3838
from flagscale.train.train import pretrain
3939
from megatron.plugin.hetero.parallel_context import get_parallel_context
4040

flagscale/train/train_qwen2_5_vl.py

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -70,8 +70,6 @@
7070
except ImportError:
7171
has_nvidia_modelopt = False
7272

73-
from flagscale.train.datasets.sft_dataset import SFTDataset, SFTDatasetConfig
74-
from flagscale.train.extra_valid import extra_valid_datasets_provider
7573
from flagscale.train.train import pretrain
7674

7775
stimer = StragglerDetector()

flagscale/train/train_qwen3_vl.py

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -53,8 +53,6 @@
5353
except ImportError:
5454
has_nvidia_modelopt = False
5555

56-
from flagscale.train.datasets.sft_dataset import SFTDatasetConfig, SFTDataset
57-
from flagscale.train.extra_valid import extra_valid_datasets_provider
5856
from flagscale.train.train import pretrain
5957
stimer = StragglerDetector()
6058

0 commit comments

Comments
 (0)