diff --git a/swift/arguments/base_args/base_args.py b/swift/arguments/base_args/base_args.py index 326fb6d8b7..1be2f8aa5a 100644 --- a/swift/arguments/base_args/base_args.py +++ b/swift/arguments/base_args/base_args.py @@ -346,6 +346,7 @@ def get_model_processor(self, res['revision'] = revision or self.model_revision res['task_type'] = task_type or self.task_type res['num_labels'] = num_labels or self.num_labels + res['processor_id_or_path'] = getattr(self, 'tokenizer_name_or_path', None) return get_model_processor(**res) diff --git a/swift/arguments/base_args/data_args.py b/swift/arguments/base_args/data_args.py index 01b8b0f7a4..a28ffa6dcb 100644 --- a/swift/arguments/base_args/data_args.py +++ b/swift/arguments/base_args/data_args.py @@ -58,6 +58,8 @@ class DataArguments: Example: '{"text1": "query", "text2": "response"}'. Defaults to None. strict (bool): If `True`, raises an error on any problematic data row. If `False`, discards the problematic sample and continues. Typically used for debugging. Defaults to False. + drop_empty_assistant_response (bool): If `True`, filters SFT rows whose assistant response is an empty string. + Defaults to False. remove_unused_columns (bool): Whether to remove columns not used by the model. If `False`, extra columns are passed to the trainer's `compute_loss` function, which is useful for custom loss calculations. Defaults to True. Note: The default is `False` for GPRO. @@ -76,6 +78,7 @@ class DataArguments: val_dataset: List[str] = field(default_factory=list) cached_dataset: List[str] = field(default_factory=list) cached_val_dataset: List[str] = field(default_factory=list) + tokenizer_name_or_path: Optional[str] = None split_dataset_ratio: float = 0. data_seed: int = 42 @@ -91,6 +94,7 @@ class DataArguments: download_mode: Literal['force_redownload', 'reuse_dataset_if_exists'] = 'reuse_dataset_if_exists' columns: Optional[Union[dict, str]] = None strict: bool = False + drop_empty_assistant_response: bool = False remove_unused_columns: bool = True disable_auto_column_mapping: bool = False # Chinese name and English name @@ -138,6 +142,7 @@ def get_dataset_kwargs(self): 'download_mode': self.download_mode, 'columns': self.columns, 'strict': self.strict, + 'drop_empty_assistant_response': self.drop_empty_assistant_response, 'model_name': self.model_name, 'model_author': self.model_author, 'remove_unused_columns': self.remove_unused_columns, diff --git a/swift/dataset/loader.py b/swift/dataset/loader.py index 11f00d48d7..6e14aab563 100644 --- a/swift/dataset/loader.py +++ b/swift/dataset/loader.py @@ -28,6 +28,7 @@ def __init__( streaming: bool = False, hub_token: Optional[str] = None, strict: bool = False, + drop_empty_assistant_response: bool = False, download_mode: Literal['force_redownload', 'reuse_dataset_if_exists'] = 'reuse_dataset_if_exists', columns: Optional[Dict[str, str]] = None, remove_unused_columns: bool = True, @@ -38,11 +39,24 @@ def __init__( self.streaming = streaming self.hub_token = hub_token self.strict = strict + self.drop_empty_assistant_response = drop_empty_assistant_response self.download_mode = download_mode self.columns = columns self.remove_unused_columns = remove_unused_columns self.disable_auto_column_mapping = disable_auto_column_mapping + @staticmethod + def _has_nonempty_assistant_response(row: Dict) -> bool: + messages = row.get('messages') or [] + return all( + message.get('role') != 'assistant' or not isinstance(message.get('content'), str) + or bool(message['content'].strip()) for message in messages) + + def _filter_empty_assistant_responses(self, dataset: HfDataset) -> HfDataset: + if not self.drop_empty_assistant_response: + return dataset + return dataset.filter(self._has_nonempty_assistant_response) + def _load_dataset_path( self, dataset_path: str, @@ -66,7 +80,7 @@ def _load_dataset_path( enable_auto_mapping=not self.disable_auto_column_mapping) if self.remove_unused_columns: dataset = RowPreprocessor.remove_useless_columns(dataset) - return dataset + return self._filter_empty_assistant_responses(dataset) def _load_repo_dataset( self, @@ -136,7 +150,7 @@ def _load_repo_dataset( enable_auto_mapping=not self.disable_auto_column_mapping) if self.remove_unused_columns: dataset = RowPreprocessor.remove_useless_columns(dataset) - datasets.append(dataset) + datasets.append(self._filter_empty_assistant_responses(dataset)) return self.concat_datasets(datasets) @staticmethod @@ -236,6 +250,7 @@ def load_dataset( use_hf: Optional[bool] = None, hub_token: Optional[str] = None, strict: bool = False, + drop_empty_assistant_response: bool = False, download_mode: Literal['force_redownload', 'reuse_dataset_if_exists'] = 'reuse_dataset_if_exists', columns: Optional[Dict[str, str]] = None, # columns_mapping remove_unused_columns: bool = True, @@ -278,6 +293,8 @@ def load_dataset( hub_token: Authentication token for accessing private datasets on the hub. Default: None. strict: If True, raise exceptions when encountering malformed data rows. If False, skip invalid rows with warnings. Default: False. + drop_empty_assistant_response: Filter rows whose assistant response is an empty string. + Defaults to False. download_mode: How to handle existing cached datasets: - 'reuse_dataset_if_exists': Use cached version if available - 'force_redownload': Always download fresh copy @@ -342,6 +359,7 @@ def load_dataset( streaming=streaming, hub_token=hub_token, strict=strict, + drop_empty_assistant_response=drop_empty_assistant_response, download_mode=download_mode, columns=columns, # columns_mapping remove_unused_columns=remove_unused_columns, diff --git a/swift/megatron/arguments/megatron_args.py b/swift/megatron/arguments/megatron_args.py index d04331f4fe..0ff44c3efa 100644 --- a/swift/megatron/arguments/megatron_args.py +++ b/swift/megatron/arguments/megatron_args.py @@ -549,6 +549,7 @@ class MegatronArguments(RLHFMegatronArgumentsMixin, MegatronTunerMixin): start_weight_decay: Optional[float] = None end_weight_decay: Optional[float] = None clip_grad: float = 1. + native_unfused_adamw: bool = False adam_beta1: float = 0.9 adam_beta2: float = 0.95 adam_eps: float = 1e-8 @@ -608,6 +609,7 @@ class MegatronArguments(RLHFMegatronArgumentsMixin, MegatronTunerMixin): overlap_param_gather: bool = False overlap_param_gather_with_optimizer_step: bool = False align_grad_reduce: bool = True + deterministic_mode: bool = False # Eagerly create NCCL communicators before the training loop to avoid the lazy # first-use allocation hitting the iteration-1 memory peak (Failed to CUDA calloc async). nccl_comm_warmup: bool = False diff --git a/swift/megatron/callbacks/print.py b/swift/megatron/callbacks/print.py index 01a41b385d..8c5e371ebd 100644 --- a/swift/megatron/callbacks/print.py +++ b/swift/megatron/callbacks/print.py @@ -1,7 +1,12 @@ # Copyright (c) ModelScope Contributors. All rights reserved. +import hashlib +import json import os +import platform +import sys import time import torch +from pathlib import Path from tqdm import tqdm from swift.megatron.utils import reduce_max_stat_across_model_parallel_group @@ -11,6 +16,145 @@ logger = get_logger() +def raw_loss_event(step, logs): + """Return an unrounded training-loss event, excluding evaluation metrics.""" + raw_losses = { + key: value + for key, value in logs.items() if key == 'loss' or (key.startswith('mtp_') and key.endswith('_loss')) + } + return {'step': step, **raw_losses} if raw_losses else None + + +def _sha256_file(path): + return hashlib.sha256(Path(path).read_bytes()).hexdigest() + + +def normalized_device(): + """Return the device class the benchmark checker expects, not the GPU model name.""" + return 'cuda' if torch.cuda.is_available() else 'cpu' + + +def normalized_dtype(value): + """Return the bench dtype alias; a bare ``torch.bfloat16`` string is rejected.""" + text = str(value or '').strip() + prefix = 'torch.' + if text.startswith(prefix): + text = text[len(prefix):] + return text.lower() + + +def machine_loss_payload(events, raw_path=None, source_sha256=None): + """Return the machine loss artifact. + + ``losses`` is the benchmark gate field: an unrounded main-loss series with one + entry per recorded step. ``events`` keeps the per-step diagnostic detail. + """ + return { + 'schema': 'glm52-machine-loss/v1', + 'framework': 'torch', + 'raw': True, + 'stage': 'training_callback_complete', + 'losses': [event['loss'] for event in events if 'loss' in event], + 'event_count': len(events), + 'steps': [event['step'] for event in events], + 'events': events, + 'source': raw_path, + 'source_sha256': source_sha256, + } + + +def _write_json(path, payload): + path = Path(path).expanduser().resolve() + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(payload, ensure_ascii=False, allow_nan=False, indent=2, sort_keys=True) + '\n') + + +def model_repro_source_modules(): + names = ( + 'swift.megatron.callbacks.print', + 'swift.megatron.trainers.base', + 'megatron.core.transformer.moe.moe_utils', + 'megatron.core.transformer.moe.experts', + 'megatron.core.transformer.moe.router', + 'megatron.core.models.common.language_module.language_module', + 'mcore_bridge.model.gpt_model', + ) + records = {} + for name in names: + path = getattr(sys.modules.get(name), '__file__', None) + records[name] = { + 'path': str(Path(path).resolve()) if path else None, + 'sha256': _sha256_file(path) if path and Path(path).is_file() else None + } + return records + + +def model_repro_topology(args, config): + from megatron.core import parallel_state + + fields = ('tensor_model_parallel_size', 'pipeline_model_parallel_size', 'expert_model_parallel_size', + 'expert_tensor_parallel_size', 'context_parallel_size', 'sequence_parallel') + record = { + 'configured': { + name: getattr(config, name, None) + for name in fields + }, + 'world_size': torch.distributed.get_world_size(), + 'rank': torch.distributed.get_rank(), + 'global_batch_size': args.global_batch_size, + 'micro_batch_size': args.micro_batch_size, + 'groups': {} + } + getters = { + 'tp': 'get_tensor_model_parallel_group', + 'pp': 'get_pipeline_model_parallel_group', + 'ep': 'get_expert_model_parallel_group', + 'etp': 'get_expert_tensor_parallel_group', + 'cp': 'get_context_parallel_group', + 'data': 'get_data_parallel_group' + } + for name, getter in getters.items(): + group = getattr(parallel_state, getter)() + record['groups'][name] = { + 'size': torch.distributed.get_world_size(group), + 'ranks': torch.distributed.get_process_group_ranks(group) + } + return record + + +def model_repro_environment(args, trainer): + """Return the formal run-local environment receipt after model loading.""" + config_path = os.environ.get('MODEL_REPRO_MODEL_CONFIG_PATH') + return { + 'schema': 'glm52-environment/v1', + 'framework': 'torch', + 'framework_version': torch.__version__, + 'python_version': platform.python_version(), + 'device': normalized_device(), + 'device_name': torch.cuda.get_device_name(torch.cuda.current_device()), + 'dtype': normalized_dtype(getattr(args, 'torch_dtype', 'bfloat16')), + 'cuda': torch.version.cuda, + 'cudnn': torch.backends.cudnn.version(), + 'nccl': list(torch.cuda.nccl.version()), + 'deterministic': { + 'algorithms_enabled': torch.are_deterministic_algorithms_enabled(), + 'cudnn_deterministic': torch.backends.cudnn.deterministic, + 'cudnn_benchmark': torch.backends.cudnn.benchmark, + 'cublas_workspace_config': os.environ.get('CUBLAS_WORKSPACE_CONFIG'), + 'nccl_algo': os.environ.get('NCCL_ALGO'), + }, + 'model_id': os.environ.get('MODEL_REPRO_MODEL_ID'), + 'revision': os.environ.get('MODEL_REPRO_MODEL_REVISION'), + 'model_config_sha256': _sha256_file(config_path) if config_path else None, + 'weights_loaded': bool(getattr(trainer, '_model_repro_weights_source', None)), + 'model_source': str(getattr(trainer, '_model_repro_weights_source', '')), + 'topology': model_repro_topology(args, trainer.config), + 'source_modules': model_repro_source_modules(), + 'invocation_id': os.environ.get('MRK_INVOCATION_ID'), + 'world_size': torch.distributed.get_world_size() if torch.distributed.is_initialized() else 1, + } + + class PrintCallback(MegatronCallback): def __init__(self, trainer): @@ -18,6 +162,8 @@ def __init__(self, trainer): self.training_bar = None self.eval_bar = None self.jsonl_writer = None + self.raw_loss_writer = None + self.raw_loss_events = [] self.is_write_rank = is_last_rank() def on_train_begin(self): @@ -27,13 +173,35 @@ def on_train_begin(self): self.training_bar.update(self.state.iteration) self.current_step = self.state.iteration self.start_time = time.time() + self.raw_loss_events = [] + loss_path = os.environ.get('MODEL_REPRO_LOSS_PATH') + if loss_path and self.is_write_rank: + Path(loss_path).expanduser().resolve().unlink(missing_ok=True) logging_path = os.path.join(self.args.output_dir, 'logging.jsonl') logger.info(f'logging_path: {logging_path}') self.jsonl_writer = JsonlWriter(logging_path, enable_async=True, write_on_rank='last') + raw_loss_path = os.environ.get('MODEL_REPRO_RAW_LOSS_PATH') + if raw_loss_path: + logger.info(f'raw_loss_path: {raw_loss_path}') + if self.is_write_rank and Path(raw_loss_path).exists(): + Path(raw_loss_path).unlink() + self.raw_loss_writer = JsonlWriter(raw_loss_path, write_on_rank='last') + env_path = os.environ.get('MODEL_REPRO_ENV_PATH') + if env_path and self.is_write_rank: + _write_json(env_path, model_repro_environment(self.args, self.trainer)) def on_train_end(self): self.training_bar.close() self.training_bar = None + loss_path = os.environ.get('MODEL_REPRO_LOSS_PATH') + if loss_path and self.is_write_rank: + raw_path = os.environ.get('MODEL_REPRO_RAW_LOSS_PATH') + payload = machine_loss_payload( + self.raw_loss_events, + raw_path=raw_path, + source_sha256=_sha256_file(raw_path) if raw_path and Path(raw_path).is_file() else None, + ) + _write_json(loss_path, payload) def on_step_end(self): n_step = self.state.iteration - self.current_step @@ -63,6 +231,11 @@ def on_log(self, logs): memory = reduce_max_stat_across_model_parallel_group(torch.cuda.max_memory_reserved() / 1024**3) logs['memory(GiB)'] = round(memory, 2) logs['train_speed(s/it)'] = round(train_speed, 6) + raw_event = raw_loss_event(state.iteration, logs) + if self.raw_loss_writer is not None and raw_event is not None: + self.raw_loss_writer.append(raw_event) + if raw_event is not None and self.is_write_rank and os.environ.get('MODEL_REPRO_LOSS_PATH'): + self.raw_loss_events.append(raw_event) logs = {k: round(v, 8) if isinstance(v, float) else v for k, v in logs.items()} self.jsonl_writer.append(logs) if self.is_write_rank: diff --git a/swift/megatron/init.py b/swift/megatron/init.py index bcfda11d7a..4a5e173af2 100644 --- a/swift/megatron/init.py +++ b/swift/megatron/init.py @@ -10,10 +10,10 @@ from copy import copy, deepcopy from packaging import version from tqdm import tqdm -from typing import Optional from transformers.modeling_utils import custom_object_save from transformers.utils import is_torch_npu_available from transformers.utils.versions import require_version +from typing import Optional from swift.model import get_model_processor, save_checkpoint from swift.utils import (HfConfigFactory, disable_safe_ddp_context_use_barrier, get_logger, get_modules_to_not_convert, @@ -22,6 +22,11 @@ logger = get_logger() +def _get_save_processor_id(args): + """Use the configured processor source when weights and tokenizer are independent.""" + return getattr(args, 'tokenizer_name_or_path', None) or args.model_dir + + def _patch__batched_p2p_ops(): from megatron.core.pipeline_parallel import p2p_communication @@ -187,6 +192,22 @@ def wrapper(*args, **kwargs): mcb_register.get_gpt_decoder_block_spec = _force_local_spec(mcb_register.get_gpt_decoder_block_spec) mcb_register.get_gpt_mtp_block_spec = _force_local_spec(mcb_register.get_gpt_mtp_block_spec) + # 1b) DSA is swapped in after the decoder spec via ModelLoader._replace_spec_dsa, + # which called _get_backend_spec_provider (TESpecProvider). That left TELinear / + # TENorm on DSA indexer + MLA while PaddleFleet HAVE_TE is False. E-259: remaining + # Torch DSA TE contaminated post_attn_norm. Force LocalSpecProvider instead. + from megatron.core.models.gpt import experimental_attention_variant_module_specs as _eav + + origin_backend_spec_provider = _eav._get_backend_spec_provider + + def _local_backend_spec_provider(config): + if not getattr(config, 'dsa_accuracy_compatible', False): + return origin_backend_spec_provider(config) + from megatron.core.models.backends import LocalSpecProvider + return LocalSpecProvider() + + _eav._get_backend_spec_provider = _local_backend_spec_provider + # 2) persist_layer_norm=False on the model config (dataclass default is baked into # __init__, so flip it on the instance via __post_init__). from mcore_bridge.config.model_config import ModelConfig as McbModelConfig @@ -221,16 +242,105 @@ def _set_layer_attn(self, mg_layer, hf_state_dict, layer_idx, to_mcore): return hf_state_dict McbGPTBridge._set_layer_attn = _set_layer_attn - logger.info('mcore_bridge patched for TE-off alignment (local spec, persist_layer_norm=False, input_layernorm map)') + + # Dense local MLPs store their norm separately from linear_fc1. Keep the + # bridge's load/export logic, correcting only this TE-specific parameter key. + origin_set_state_dict = McbGPTBridge._set_state_dict + + def _set_state_dict(self, mg_module, mg_key, hf_state_dict, hf_key, to_mcore, **kwargs): + if mg_key == 'mlp.linear_fc1.layer_norm_weight': + fc1 = getattr(getattr(mg_module, 'mlp', None), 'linear_fc1', None) + if getattr(fc1, 'layer_norm_weight', None) is None: + mg_key = 'pre_mlp_layernorm.weight' + return origin_set_state_dict(self, mg_module, mg_key, hf_state_dict, hf_key, to_mcore, **kwargs) + + McbGPTBridge._set_state_dict = _set_state_dict + logger.info( + 'mcore_bridge patched for TE-off alignment (local spec, persist_layer_norm=False, input_layernorm+mlp-norm map)' + ) + + +def _patch_mcore_bridge_tp1_accuracy(): + """Apply the DSA TP1 graph choice only within the configured bridge instance.""" + from contextvars import ContextVar + from functools import wraps + from mcore_bridge.model.modules import mtp_layer, transformer_block + + active = ContextVar('swift_dsa_tp1_accuracy', default=False) + + def scoped(method): + + @wraps(method) + def call(self, *args, **kwargs): + config = self.config + token = active.set( + getattr(config, 'dsa_accuracy_compatible', False) and config.tensor_model_parallel_size <= 1) + try: + return method(self, *args, **kwargs) + finally: + active.reset(token) + + return call + + def patch_module(module, cls, methods): + original = module.make_viewless_tensor + if getattr(original, '_swift_tp1_accuracy_patch', False): + return + + def make_viewless_tensor(inp, requires_grad, keep_graph): + if active.get(): + return inp + return original(inp=inp, requires_grad=requires_grad, keep_graph=keep_graph) + + make_viewless_tensor._swift_tp1_accuracy_patch = True + module.make_viewless_tensor = make_viewless_tensor + if hasattr(module, 'gather_from_tensor_model_parallel_region'): + original_gather = module.gather_from_tensor_model_parallel_region + + def gather(input_, group=None): + if active.get() and (group is None or group.size() <= 1): + return input_ + return original_gather(input_, group=group) + + module.gather_from_tensor_model_parallel_region = gather + for name in methods: + setattr(cls, name, scoped(getattr(cls, name))) + + patch_module(mtp_layer, mtp_layer.MultiTokenPredictionLayer, ('_concat_embeddings', '_get_embeddings')) + patch_module(transformer_block, transformer_block.TransformerBlock, ('forward', )) def _patch_mcore_bridge(): require_version('mcore-bridge>=1.4.0', 'please install mcore-bridge via `pip install mcore-bridge -U`') import mcore_bridge from mcore_bridge import GPTBridge + from mcore_bridge.model.register import ModelLoader logger.info(f'mcore_bridge.__version__: {mcore_bridge.__version__}') if _use_accuracy_compatible_enabled(): _patch_mcore_bridge_disable_te() + _patch_mcore_bridge_tp1_accuracy() + if not getattr(ModelLoader._replace_spec_dsa, '_swift_norm_accuracy_patch', False): + origin_replace_spec_dsa = ModelLoader._replace_spec_dsa + + def replace_spec_dsa(self, layer_spec): + origin_replace_spec_dsa(self, layer_spec) + from megatron.core.transformer.torch_norm import WrappedTorchNorm + + dsa_spec = layer_spec.submodules.self_attention + if getattr(self.config, 'norm_accuracy_compatible', False): + dsa_spec.submodules.q_layernorm = WrappedTorchNorm + dsa_spec.submodules.kv_layernorm = WrappedTorchNorm + indexer = getattr( + getattr(dsa_spec.submodules.core_attention, 'submodules', None), + 'indexer', + None, + ) + if (getattr(self.config, 'norm_accuracy_compatible', False) and indexer is not None + and getattr(indexer, 'submodules', None) is not None): + indexer.submodules.k_norm = WrappedTorchNorm + + replace_spec_dsa._swift_norm_accuracy_patch = True + ModelLoader._replace_spec_dsa = replace_spec_dsa origin_save_weights = GPTBridge.save_weights def save_weights( @@ -255,7 +365,11 @@ def save_weights( else: with torch.device('meta'), disable_safe_ddp_context_use_barrier(): self.hf_model = get_model_processor( - args.model_dir, model_type=args.model_type, return_dummy_model=True)[0] + args.model_dir, + model_type=args.model_type, + return_dummy_model=True, + processor_id_or_path=_get_save_processor_id(args), + )[0] if is_master(): if peft_format: diff --git a/swift/megatron/model/utils.py b/swift/megatron/model/utils.py index acc2a7e48b..96cac5eb78 100644 --- a/swift/megatron/model/utils.py +++ b/swift/megatron/model/utils.py @@ -5,7 +5,7 @@ from mcore_bridge import hf_to_mcore_config from transformers.utils import is_torch_npu_available -from swift.utils import get_logger +from swift.utils import HfConfigFactory, get_logger logger = get_logger() @@ -35,8 +35,25 @@ def _check_padding_free(args, config): args.padding_free = False +def _check_dsa_index_share_recompute(config): + """Reject activation replay that omits a DSA skip layer's source indexer.""" + if (config.experimental_attention_variant == 'dsa' and (getattr(config, 'dsa_indexer_topk_freq', 1) or 1) > 1 + and getattr(config, 'recompute_granularity', None) not in {None, 'none'}): + raise ValueError( + 'DSA cross-layer top-k sharing is incompatible with activation recompute because a skip layer may be ' + 'replayed without its source computing layer. Set recompute_granularity=none.') + + def get_mcore_model_config(args, hf_config): kwargs = hf_to_mcore_config(hf_config) + llm_config = HfConfigFactory.get_text_config(hf_config) + n_routed_experts = getattr(llm_config, 'n_routed_experts', None) + if getattr(llm_config, 'model_type', None) == 'glm_moe_dsa': + kwargs['accuracy_compatible_loss_sum_dtype'] = 'float32' + if n_routed_experts is not None: + kwargs['num_moe_experts'] = n_routed_experts + # Checkpoint MTP metadata describes available weights, not an opt-in to + # auxiliary training. The explicit mtp_num_layers argument below controls it. kwargs['mcore_model_type'] = args.megatron_model_meta.model_type kwargs['hf_config'] = hf_config for f in fields(ModelConfig): @@ -72,11 +89,14 @@ def get_mcore_model_config(args, hf_config): kwargs['moe_enable_routing_replay'] = True if args.megatron_extra_kwargs: kwargs.update(args.megatron_extra_kwargs) + if kwargs.get('accuracy_compatible_loss_sum_dtype', 'float64') not in {'float32', 'float64'}: + raise ValueError('accuracy_compatible_loss_sum_dtype must be float32 or float64') config = ModelConfig(**kwargs) if is_torch_npu_available() and getattr(args, 'attention_backend', 'flash') != 'local': setattr(config, 'use_flash_attn', True) _check_attention_backend(args, config) _check_padding_free(args, config) + _check_dsa_index_share_recompute(config) return config diff --git a/swift/megatron/trainers/base.py b/swift/megatron/trainers/base.py index a4e4349f5b..0b0d5da9ad 100644 --- a/swift/megatron/trainers/base.py +++ b/swift/megatron/trainers/base.py @@ -116,6 +116,7 @@ def _load_checkpoint(self): if args.mcore_model is not None: self.state.iteration = load_mcore_checkpoint( args, self.wrapped_models, self.optimizer, self.opt_param_scheduler, load_arg='mcore_model') + self._model_repro_weights_source = args.mcore_model if args.mcore_adapter is not None: self.state.iteration = load_mcore_checkpoint( args, self.wrapped_models, self.optimizer, self.opt_param_scheduler, load_arg='mcore_adapter') @@ -199,6 +200,7 @@ def _prepare_peft_model(self, models): args = self.args if args.mcore_model is None: self.bridge.load_weights(models, args.model_dir) + self._model_repro_weights_source = args.model_dir peft_models = [prepare_mcore_model(args, model) for model in models] if args.tuner_type == 'lora' and args.adapters and args.mcore_adapter is None: assert len(args.adapters) == 1, 'Currently only support one adapter.' @@ -776,9 +778,13 @@ def save_checkpoint(self): state = self.state args.consumed_train_samples = state.consumed_train_samples iteration = state.iteration - output_dir = os.path.join(args.output_dir, f'checkpoint-{iteration}') + formal_checkpoint_dir = os.environ.get('MODEL_REPRO_CHECKPOINT_DIR') + if formal_checkpoint_dir and iteration == args.train_iters: + output_dir = os.path.abspath(os.path.expanduser(formal_checkpoint_dir)) + else: + output_dir = os.path.join(args.output_dir, f'checkpoint-{iteration}') os.makedirs(output_dir, exist_ok=True) - args_path = os.path.join(os.path.dirname(output_dir), 'args.json') + args_path = os.path.join(args.output_dir, 'args.json') self.copy_path(args_path, os.path.join(output_dir, 'args.json')) if args.save_safetensors and args.no_save_optim: model = [] diff --git a/swift/megatron/trainers/trainer.py b/swift/megatron/trainers/trainer.py index 46fa9914d4..603fafead1 100644 --- a/swift/megatron/trainers/trainer.py +++ b/swift/megatron/trainers/trainer.py @@ -1,4 +1,7 @@ # Copyright (c) ModelScope Contributors. All rights reserved. +import hashlib +import json +import os import torch import torch.distributed as dist import torch.nn @@ -16,8 +19,101 @@ logger = get_logger() +def project_owning_loader_semantics(input_values, model_label_values, semantic_length, labels_were_shifted=True): + """Normalize padded Megatron carrier tensors back to the dataset semantic row.""" + semantic_length = int(semantic_length) + if semantic_length <= 0 or semantic_length > len(input_values) or semantic_length > len(model_label_values): + raise ValueError(f'invalid owning-loader semantic length {semantic_length} for carrier lengths ' + f'{len(input_values)}/{len(model_label_values)}') + semantic_input_values = input_values[:semantic_length] + normalized_label_values = model_label_values + if labels_were_shifted and model_label_values: + # get_batch_on_this_pp_rank rolls causal-LM labels left by one before + # model forward. Reverse that roll for a framework-neutral dataset receipt. + normalized_label_values = model_label_values[-1:] + model_label_values[:-1] + semantic_label_values = normalized_label_values[:semantic_length] + semantic_mask_values = [label != -100 for label in semantic_label_values] + return semantic_input_values, semantic_label_values, semantic_mask_values + + class MegatronTrainer(BaseMegatronTrainer): + def _write_input_contract_once(self, data, seq_lens=None): + path = os.environ.get('MODEL_REPRO_INPUT_RECEIPT_PATH') + if not path or getattr(self, '_input_contract_written', False): + return + if not mpu.is_pipeline_last_stage(ignore_virtual=False): + return + if (torch.distributed.is_initialized() + and torch.distributed.get_rank() != torch.distributed.get_world_size() - 1): + return + input_ids = data.get('input_ids') + labels = data.get('labels') + if input_ids is None or labels is None: + return + + def values(tensor): + return tensor.detach().to(device='cpu', dtype=torch.int64).reshape(-1).tolist() + + def digest(items): + return hashlib.sha256(json.dumps(items, separators=(',', ':')).encode()).hexdigest() + + input_values = values(input_ids) + label_values = values(labels) + model_mask_values = [label != -100 for label in label_values] + semantic_length = seq_lens[0] if seq_lens else len(input_values) + labels_were_shifted = self.args.task_type == 'causal_lm' + semantic_input_values, semantic_label_values, semantic_mask_values = project_owning_loader_semantics( + input_values, label_values, semantic_length, labels_were_shifted) + payload = { + 'schema': 'glm52-owning-loader-input/v1', + 'framework': 'torch', + 'rank': torch.distributed.get_rank() if torch.distributed.is_initialized() else 0, + 'step': self.state.iteration + 1, + 'input_ids': { + 'shape': list(input_ids.shape), + 'dtype': str(input_ids.dtype), + 'count': len(input_values), + 'sha256': digest(input_values), + }, + 'labels': { + 'shape': list(labels.shape), + 'dtype': str(labels.dtype), + 'count': len(label_values), + 'supervised_count': sum(model_mask_values), + 'sha256': digest(label_values), + 'projection': 'model_next_token_labels', + }, + 'loss_mask': { + 'shape': list(labels.shape), + 'dtype': 'bool', + 'count': len(model_mask_values), + 'supervised_count': sum(model_mask_values), + 'sha256': digest(model_mask_values), + }, + 'semantic': { + 'input_token_count': len(semantic_input_values), + 'supervised_target_count': sum(semantic_mask_values), + 'input_ids_sha256': digest(semantic_input_values), + 'labels_sha256': digest(semantic_label_values), + 'loss_mask_sha256': digest(semantic_mask_values), + 'projection': 'dataset_row_before_megatron_padding_and_label_roll', + }, + 'carrier_padding': { + 'count': len(input_values) - len(semantic_input_values), + 'input_ids_sha256': digest(input_values[len(semantic_input_values):]), + 'labels_sha256': digest(label_values[len(semantic_input_values):]), + }, + 'ignore_index': -100, + 'dataset': os.environ.get('MODEL_REPRO_INPUT_DATASET_PATH'), + } + path = os.path.abspath(os.path.expanduser(path)) + os.makedirs(os.path.dirname(path), exist_ok=True) + with open(path, 'w', encoding='utf-8') as stream: + json.dump(payload, stream, ensure_ascii=False, indent=2, sort_keys=True) + stream.write('\n') + self._input_contract_written = True + def seq_cls_loss_func(self, output_tensor, *, labels: torch.Tensor, packed_seq_params=None, attention_mask=None): args = self.args logits = self.get_last_tokens(output_tensor, packed_seq_params, attention_mask) @@ -61,11 +157,11 @@ def loss_func(self, losses = losses * torch.exp(-losses.detach()) if loss_scale is not None: losses = losses * loss_scale - from megatron.core.transformer.module import _use_accuracy_compatible - if _use_accuracy_compatible(): - loss_sum = (losses * loss_mask).reshape(-1).double().sum().float() + masked_losses = losses * loss_mask + if _use_accuracy_compatible_enabled() and self.config.accuracy_compatible_loss_sum_dtype == 'float64': + loss_sum = masked_losses.reshape(-1).double().sum().float() else: - loss_sum = torch.sum(losses * loss_mask) + loss_sum = torch.sum(masked_losses) loss = torch.cat([loss_sum.view(1), loss_mask.sum().view(1)]) # Reduce loss for logging. @@ -125,6 +221,8 @@ def _compute_channel_loss(self, losses, loss_mask, channels, packed_seq_params=N def forward_step(self, data_iterator, model): vp_stage = model.module.module.vp_stage data = self.get_batch(data_iterator, vp_stage) + seq_lens = data.pop('_model_repro_seq_lens', None) + self._write_input_contract_once(data, seq_lens) loss_scale = data.pop('loss_scale', None) channels = data.pop('channel', None) labels = data.get('labels') diff --git a/swift/megatron/trainers/utils.py b/swift/megatron/trainers/utils.py index 16ca70d9a2..a77d4e6a4f 100644 --- a/swift/megatron/trainers/utils.py +++ b/swift/megatron/trainers/utils.py @@ -1,5 +1,6 @@ # Copyright (c) ModelScope Contributors. All rights reserved. import gc +import os import torch from accelerate.utils import gather as hf_gather from accelerate.utils import gather_object as hf_gather_object @@ -400,6 +401,10 @@ def prepare_batch(args, data, vp_stage=None): if num_samples is not None: batch['packed_seq_params'].num_samples = num_samples batch = get_batch_on_this_cp_rank(args, batch) + if os.environ.get('MODEL_REPRO_INPUT_RECEIPT_PATH') and seq_lens is not None: + # Opt-in metadata for the owning-loader receipt. MegatronTrainer removes + # it before model(**data), so it cannot alter the numerical path. + batch['_model_repro_seq_lens'] = [int(length) for length in seq_lens] return batch diff --git a/swift/megatron/utils/megatron_lm_utils.py b/swift/megatron/utils/megatron_lm_utils.py index 55e4c1d80f..82ac657b2d 100644 --- a/swift/megatron/utils/megatron_lm_utils.py +++ b/swift/megatron/utils/megatron_lm_utils.py @@ -83,7 +83,30 @@ def _initialize_mpu(args): f'EP: {args.expert_model_parallel_size}, ETP: {args.expert_tensor_parallel_size}') +def configure_deterministic_mode(args): + """Enable the same fail-closed deterministic contract as Megatron-LM training.""" + if not getattr(args, 'deterministic_mode', False): + return + attention_backend = getattr(args, 'attention_backend', '') + attention_backend = getattr(attention_backend, 'name', str(attention_backend)).lower() + if attention_backend == 'flash': + raise ValueError('Flash attention cannot be used in deterministic mode.') + if getattr(args, 'cross_entropy_loss_fusion', False): + raise ValueError('Cross entropy fusion cannot be used in deterministic mode.') + nccl_algo = os.environ.get('NCCL_ALGO') + allowed_nccl_algorithms = {'Tree', 'Ring', 'CollnetDirect', 'CollnetChain', '^NVLS'} + if nccl_algo not in allowed_nccl_algorithms: + raise ValueError(f'NCCL_ALGO must be explicitly set to one of {sorted(allowed_nccl_algorithms)} ' + 'in deterministic mode.') + torch.use_deterministic_algorithms(True) + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + logger.info(f'Deterministic mode enabled with NCCL_ALGO={nccl_algo}.') + + def initialize_megatron(args): + configure_deterministic_mode(args) + # Pytorch distributed. _initialize_mpu(args) diff --git a/swift/megatron/utils/utils.py b/swift/megatron/utils/utils.py index 9562475849..4774c99095 100644 --- a/swift/megatron/utils/utils.py +++ b/swift/megatron/utils/utils.py @@ -212,6 +212,9 @@ def get_padding_to(args): padding_to = None if args.tensor_model_parallel_size > 1 and args.sequence_parallel: padding_to = args.tensor_model_parallel_size + # Match the DSA reference carrier without changing other TP+SP models. + if (getattr(args, 'megatron_extra_kwargs', None) or {}).get('dsa_accuracy_compatible', False): + padding_to *= 2 if args.context_parallel_size > 1: padding_to = (padding_to or 1) * args.context_parallel_size origin_padding_to = padding_to diff --git a/swift/model/register.py b/swift/model/register.py index e106929d13..842b94e5f6 100644 --- a/swift/model/register.py +++ b/swift/model/register.py @@ -174,6 +174,7 @@ def __init__( auto_model_cls=None, return_dummy_model: bool = False, new_special_tokens: Optional[List[str]] = None, + processor_id_or_path: Optional[str] = None, model_kwargs: Optional[Dict[str, Any]] = None, **kwargs, ): @@ -197,6 +198,7 @@ def __init__( self.auto_tokenizer_cls = None self.return_dummy_model = return_dummy_model self.new_special_tokens = new_special_tokens + self.processor_id_or_path = processor_id_or_path self.model_kwargs = model_kwargs self.patch_offload = kwargs.pop('patch_offload', False) self.init_strategy = kwargs.get('init_strategy') @@ -257,15 +259,16 @@ def _get_tokenizer(self, processor): return tokenizer def get_processor(self, model_dir: str, config: PretrainedConfig) -> Processor: + processor_dir = self.processor_id_or_path or model_dir auto_tokenizer_cls = self.auto_tokenizer_cls if auto_tokenizer_cls is None: - if os.path.exists(os.path.join(model_dir, 'preprocessor_config.json')) or os.path.exists( - os.path.join(model_dir, 'processor_config.json')): + if os.path.exists(os.path.join(processor_dir, 'preprocessor_config.json')) or os.path.exists( + os.path.join(processor_dir, 'processor_config.json')): from transformers import AutoProcessor auto_tokenizer_cls = AutoProcessor else: auto_tokenizer_cls = AutoTokenizer - return auto_tokenizer_cls.from_pretrained(model_dir, trust_remote_code=self.default_trust_remote_code) + return auto_tokenizer_cls.from_pretrained(processor_dir, trust_remote_code=self.default_trust_remote_code) def get_model(self, model_dir: str, config: PretrainedConfig, processor: Processor, model_kwargs) -> PreTrainedModel: @@ -534,6 +537,7 @@ def get_model_processor( max_model_len: Optional[int] = None, auto_model_cls=None, new_special_tokens: Optional[List[str]] = None, + processor_id_or_path: Optional[str] = None, task_type: Literal['causal_lm', 'seq_cls', 'embedding', 'reranker', 'generative_reranker'] = None, num_labels: Optional[int] = None, problem_type: Literal['regression', 'single_label_classification', 'multi_label_classification'] = None, @@ -568,6 +572,7 @@ def get_model_processor( max_model_len: Maximum sequence length the model can handle. auto_model_cls: Custom AutoModel class to use for loading (e.g., AutoModelForCausalLM). new_special_tokens: List of new special tokens to add to the tokenizer. + processor_id_or_path: Optional independent tokenizer or processor source. task_type: Task type for the model. Options: 'causal_lm', 'seq_cls', 'embedding', 'reranker', 'generative_reranker'. num_labels: Number of labels for classification tasks. @@ -625,6 +630,7 @@ def get_model_processor( auto_model_cls=auto_model_cls, return_dummy_model=return_dummy_model, new_special_tokens=new_special_tokens, + processor_id_or_path=processor_id_or_path, model_kwargs=model_kwargs, **kwargs) return loader.load() diff --git a/tests/general/test_dataset_empty_assistant.py b/tests/general/test_dataset_empty_assistant.py new file mode 100644 index 0000000000..2132de2bb0 --- /dev/null +++ b/tests/general/test_dataset_empty_assistant.py @@ -0,0 +1,67 @@ +from datasets import Dataset + +from swift.dataset.loader import DatasetLoader + + +def test_empty_assistant_filter_is_default_off(): + dataset = Dataset.from_list([ + { + 'messages': [{ + 'role': 'user', + 'content': 'q' + }, { + 'role': 'assistant', + 'content': '' + }] + }, + { + 'messages': [{ + 'role': 'user', + 'content': 'q' + }, { + 'role': 'assistant', + 'content': 'a' + }] + }, + ]) + + filtered = DatasetLoader()._filter_empty_assistant_responses(dataset) + + assert len(filtered) == 2 + + +def test_empty_assistant_filter_drops_only_blank_assistant_content(): + dataset = Dataset.from_list([ + { + 'messages': [{ + 'role': 'user', + 'content': 'q' + }, { + 'role': 'assistant', + 'content': '' + }] + }, + { + 'messages': [{ + 'role': 'user', + 'content': 'q' + }, { + 'role': 'assistant', + 'content': ' ' + }] + }, + { + 'messages': [{ + 'role': 'user', + 'content': '' + }, { + 'role': 'assistant', + 'content': 'a' + }] + }, + ]) + + filtered = DatasetLoader(drop_empty_assistant_response=True)._filter_empty_assistant_responses(dataset) + + assert len(filtered) == 1 + assert filtered[0]['messages'][-1]['content'] == 'a' diff --git a/tests/megatron/test_accuracy_bridge_tp1.py b/tests/megatron/test_accuracy_bridge_tp1.py new file mode 100644 index 0000000000..9c8ed89768 --- /dev/null +++ b/tests/megatron/test_accuracy_bridge_tp1.py @@ -0,0 +1,104 @@ +"""Instance-scoped TP1 bridge behavior, including mixed configurations.""" +import ast +import sys +import unittest +from pathlib import Path +from types import ModuleType, SimpleNamespace +from unittest.mock import patch + + +def load_patch(): + path = Path(__file__).resolve().parents[2] / 'swift/megatron/init.py' + node = next( + node for node in ast.parse(path.read_text()).body + if isinstance(node, ast.FunctionDef) and node.name == '_patch_mcore_bridge_tp1_accuracy') + namespace = {} + exec(compile(ast.Module(body=[node], type_ignores=[]), str(path), 'exec'), namespace) + return namespace[node.name] + + +class TestBridgeTp1Patch(unittest.TestCase): + + def setUp(self): + self.modules = [] + for name, class_name in (('mtp_layer', 'MultiTokenPredictionLayer'), ('transformer_block', 'TransformerBlock')): + module = ModuleType(name) + module.make_viewless_tensor = lambda inp, **kwargs: ('wrapped', inp) + module.gather_from_tensor_model_parallel_region = lambda inp, **kwargs: ('gathered', inp) + + def forward(self, inp, callback=None, module=module): + result = module.make_viewless_tensor(inp=inp, requires_grad=True, keep_graph=True) + if callback: + callback() + return result + + cls = type(class_name, (), {'forward': forward, '_concat_embeddings': forward, '_get_embeddings': forward}) + setattr(module, class_name, cls) + self.modules.append((module, cls)) + modules = ModuleType('mcore_bridge.model.modules') + modules.mtp_layer, modules.transformer_block = [item[0] for item in self.modules] + replacement = patch.dict(sys.modules, {'mcore_bridge.model.modules': modules}) + replacement.start() + self.addCleanup(replacement.stop) + self.apply_patch = load_patch() + self.apply_patch() + + def instance(self, cls, enabled, tp_size=1): + instance = cls() + instance.config = SimpleNamespace(dsa_accuracy_compatible=enabled, tensor_model_parallel_size=tp_size) + return instance + + def test_idempotent(self): + originals = [(module.make_viewless_tensor, cls.forward) for module, cls in self.modules] + self.apply_patch() + for (module, cls), (function, method) in zip(self.modules, originals): + self.assertIs(module.make_viewless_tensor, function) + self.assertIs(cls.forward, method) + + def test_only_explicit_dsa_tp1_instance_skips_viewless(self): + inp = object() + for module, cls in self.modules: + name = '_concat_embeddings' if module.__name__ == 'mtp_layer' else 'forward' + for enabled, tp_size in ((False, 1), (True, 1), (True, 2)): + output = getattr(self.instance(cls, enabled, tp_size), name)(inp) + if enabled and tp_size == 1: + self.assertIs(output, inp) + else: + self.assertEqual(output, ('wrapped', inp)) + self.assertEqual(module.make_viewless_tensor(inp, True, True), ('wrapped', inp)) + + def test_gather_uses_instance_scope_and_actual_group(self): + module, cls = self.modules[1] + inp = object() + for enabled, tp_size in ((False, 1), (True, 1), (True, 2)): + observed = [] + group = SimpleNamespace(size=lambda: tp_size) + + def callback(): + observed.append(module.gather_from_tensor_model_parallel_region(inp, group)) + + self.instance(cls, enabled).forward(inp, callback=callback) + if enabled and tp_size == 1: + self.assertIs(observed[0], inp) + else: + self.assertEqual(observed[0], ('gathered', inp)) + self.assertEqual(module.gather_from_tensor_model_parallel_region(inp), ('gathered', inp)) + + def test_nested_legacy_call_and_exception_restore_scope(self): + module, cls = self.modules[1] + enabled = self.instance(cls, True) + legacy = self.instance(cls, False) + inp = object() + + def nested(): + self.assertEqual(legacy.forward(inp), ('wrapped', inp)) + self.assertIs(module.make_viewless_tensor(inp, True, True), inp) + raise ValueError('test failure') + + with self.assertRaisesRegex(ValueError, 'test failure'): + enabled.forward(inp, callback=nested) + self.assertEqual(module.make_viewless_tensor(inp, True, True), ('wrapped', inp)) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/megatron/test_accuracy_loss_and_norm.py b/tests/megatron/test_accuracy_loss_and_norm.py new file mode 100644 index 0000000000..235f94ab32 --- /dev/null +++ b/tests/megatron/test_accuracy_loss_and_norm.py @@ -0,0 +1,110 @@ +"""Native criterion and DSA specification contracts without distributed setup.""" + +import ast +import contextlib +import io +import sys +import torch +import types +import unittest +from pathlib import Path +from unittest.mock import patch + +from swift.megatron.trainers import trainer as trainer_module + +ROOT = Path(__file__).resolve().parents[2] + + +def production_function(relative_path, name, namespace): + path = ROOT / relative_path + tree = ast.parse(path.read_text()) + matches = [node for node in ast.walk(tree) if isinstance(node, ast.FunctionDef) and node.name == name] + if len(matches) != 1: + raise AssertionError(f'Expected one production function: {path}:{name}') + node = matches[0] + node.decorator_list = [] + module = ast.Module( + body=[ast.ImportFrom(module='__future__', names=[ast.alias(name='annotations')], level=0), node], + type_ignores=[]) + exec(compile(ast.fix_missing_locations(module), str(path), 'exec'), namespace) + return namespace[name] + + +class AccuracyLossAndNormTest(unittest.TestCase): + + def test_mask_scale_local_gradient_and_global_reporting(self): + torch.cuda.set_device(0) + for enabled in (False, True): + with self.subTest(accuracy=enabled): + loss_func = trainer_module.MegatronTrainer.loss_func + trainer = types.SimpleNamespace( + args=types.SimpleNamespace(enable_dft_loss=False, enable_channel_loss=False), + config=types.SimpleNamespace(accuracy_compatible_loss_sum_dtype='float64')) + values = torch.tensor([[2., 19., 3.]], device='cuda', requires_grad=True) + labels = torch.tensor([[1, -100, 2]], device='cuda') + scale = torch.tensor([[0.5, 1000., 2.]], device='cuda') + # A second identical DP rank contributes to reporting, not local backward. + with patch.object(torch.distributed, 'all_reduce', side_effect=lambda value, **kwargs: value.mul_(2)), \ + patch.object(torch.distributed, 'get_rank', return_value=0), \ + patch.object(trainer_module.mpu, 'get_data_parallel_group', return_value=None), \ + patch.object(trainer_module, '_use_accuracy_compatible_enabled', return_value=enabled), \ + contextlib.redirect_stdout(io.StringIO()): + loss, count, metrics = loss_func(trainer, values, labels=labels, loss_scale=scale) + self.assertEqual(loss.item(), 7.) + self.assertEqual(count.item(), 2) + self.assertEqual(metrics['loss'].tolist(), [14., 4.]) + self.assertFalse(metrics['loss'].requires_grad) + loss.backward() + self.assertEqual(values.grad.tolist(), [[0.5, 0., 2.]]) + + def test_fp64_compatibility_sum_keeps_small_losses_and_masked_gradients(self): + torch.cuda.set_device(0) + trainer = types.SimpleNamespace( + args=types.SimpleNamespace(enable_dft_loss=False, enable_channel_loss=False), + config=types.SimpleNamespace(accuracy_compatible_loss_sum_dtype='float64')) + values = torch.tensor([[100000000.] + [1.] * 8 + [100000000.]], device='cuda', requires_grad=True) + labels = torch.tensor([[1] * 9 + [-100]], device='cuda') + with patch.object(torch.distributed, 'all_reduce'), \ + patch.object(torch.distributed, 'get_rank', return_value=0), \ + patch.object(trainer_module.mpu, 'get_data_parallel_group', return_value=None), \ + patch.object(trainer_module, '_use_accuracy_compatible_enabled', return_value=True), \ + contextlib.redirect_stdout(io.StringIO()): + loss, count, metrics = trainer_module.MegatronTrainer.loss_func(trainer, values, labels=labels) + self.assertEqual(loss.item(), 100000008.) + self.assertEqual(count.item(), 9) + self.assertEqual(metrics['loss'].tolist(), [100000008., 9.]) + loss.backward() + self.assertEqual(values.grad.tolist(), [[1.] * 9 + [0.]]) + + def test_indexer_norm_preserves_disabled_provider(self): + native_norm = type('NativeNorm', (), {}) + provider_norm = type('ProviderNorm', (), {}) + module = types.ModuleType('megatron.core.transformer.torch_norm') + module.WrappedTorchNorm = native_norm + for enabled, norm_accuracy in ((False, False), (True, False), (True, True)): + with self.subTest(accuracy=enabled, norm_accuracy=norm_accuracy): + calls = [] + replace_spec = production_function( + 'swift/megatron/init.py', 'replace_spec_dsa', { + 'origin_replace_spec_dsa': lambda *args: calls.append('provider'), + '_use_accuracy_compatible_enabled': lambda: enabled, + }) + indexer = types.SimpleNamespace(submodules=types.SimpleNamespace(k_norm=provider_norm)) + attention = types.SimpleNamespace( + submodules=types.SimpleNamespace( + q_layernorm=provider_norm, + kv_layernorm=provider_norm, + core_attention=types.SimpleNamespace(submodules=types.SimpleNamespace(indexer=indexer)))) + spec = types.SimpleNamespace(submodules=types.SimpleNamespace(self_attention=attention)) + loader = types.SimpleNamespace(config=types.SimpleNamespace(norm_accuracy_compatible=norm_accuracy)) + with patch.dict(sys.modules, {module.__name__: module}): + replace_spec(loader, spec) + self.assertEqual(calls, ['provider']) + self.assertIs(indexer.submodules.k_norm, native_norm if norm_accuracy else provider_norm) + expected_qkv = native_norm if norm_accuracy else provider_norm + self.assertIs(attention.submodules.q_layernorm, expected_qkv) + self.assertIs(attention.submodules.kv_layernorm, expected_qkv) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/megatron/test_model_config.py b/tests/megatron/test_model_config.py new file mode 100644 index 0000000000..bbf85f99d0 --- /dev/null +++ b/tests/megatron/test_model_config.py @@ -0,0 +1,200 @@ +import ast +import inspect +import math +import torch +from transformers import PretrainedConfig +from types import SimpleNamespace + +from swift.megatron.init import _get_save_processor_id, _patch_mcore_bridge_disable_te +from swift.megatron.model import utils +from swift.megatron.utils.utils import get_padding_to + + +class _ModelConfigStub: + + def __init__(self, **kwargs): + self.kwargs = kwargs + self.attention_backend = SimpleNamespace(name='unfused') + self.experimental_attention_variant = 'dsa' + + +def _make_args(mtp_num_layers=None): + return SimpleNamespace( + megatron_model_meta=SimpleNamespace(model_type='gpt'), + mtp_num_layers=mtp_num_layers, + task_type='causal_lm', + torch_dtype=torch.bfloat16, + decoder_first_pipeline_num_layers=None, + decoder_last_pipeline_num_layers=None, + fp4_param_gather=False, + fp8_param_gather=False, + moe_grouped_gemm=False, + router_replay_mode='disabled', + megatron_extra_kwargs=None, + padding_free=False, + ) + + +def _patch_model_config(monkeypatch): + monkeypatch.setattr(utils, 'ModelConfig', _ModelConfigStub) + monkeypatch.setattr(utils, 'fields', + lambda _: [SimpleNamespace(name='mtp_num_layers'), + SimpleNamespace(name='num_moe_experts')]) + + +def test_save_processor_prefers_independent_tokenizer_source(): + assert _get_save_processor_id(SimpleNamespace(model_dir='/weights', + tokenizer_name_or_path='/tokenizer')) == '/tokenizer' + assert _get_save_processor_id(SimpleNamespace(model_dir='/weights', tokenizer_name_or_path=None)) == '/weights' + + +def test_get_mcore_model_config_propagates_accuracy_mode(monkeypatch): + # Preserve the real inherited dataclass fields while avoiding model construction. + actual_fields = utils.fields(utils.ModelConfig) + assert 'use_accuracy_compatible' in {field.name for field in actual_fields} + monkeypatch.setattr(utils, 'ModelConfig', _ModelConfigStub) + monkeypatch.setattr(utils, 'fields', lambda _: actual_fields) + for enabled in (False, True, False): + args = _make_args() + args.use_accuracy_compatible = enabled + monkeypatch.setenv('USE_ACCURACY_COMPATIBLE', str(int(not enabled))) + config = utils.get_mcore_model_config(args, PretrainedConfig()) + assert config.kwargs['use_accuracy_compatible'] is enabled + + +def test_get_mcore_model_config_does_not_enable_mtp_from_checkpoint(monkeypatch): + _patch_model_config(monkeypatch) + hf_config = PretrainedConfig(num_nextn_predict_layers=1) + + config = utils.get_mcore_model_config(_make_args(), hf_config) + + assert not config.kwargs.get('mtp_num_layers') + + +def test_glm52_loss_sum_contract_keeps_other_models_default(monkeypatch): + _patch_model_config(monkeypatch) + for model_type in ('glm_moe_dsa', 'glm4_moe', 'minimax_m2'): + hf_config = PretrainedConfig(model_type=model_type) + config = utils.get_mcore_model_config(_make_args(), hf_config) + assert config.kwargs.get('accuracy_compatible_loss_sum_dtype', + 'float64') == ('float32' if model_type == 'glm_moe_dsa' else 'float64') + + +def test_loss_sum_contract_rejects_unsupported_dtype(monkeypatch): + _patch_model_config(monkeypatch) + args = _make_args() + args.megatron_extra_kwargs = {'accuracy_compatible_loss_sum_dtype': 'bfloat16'} + try: + utils.get_mcore_model_config(args, PretrainedConfig()) + except ValueError as error: + assert 'accuracy_compatible_loss_sum_dtype' in str(error) + else: + raise AssertionError('BF16 loss accumulation must fail before constructing the model') + + +def test_get_mcore_model_config_does_not_enable_mtp_from_nested_checkpoint(monkeypatch): + _patch_model_config(monkeypatch) + hf_config = PretrainedConfig(text_config=PretrainedConfig(mtp_num_hidden_layers=1)) + + config = utils.get_mcore_model_config(_make_args(), hf_config) + + assert not config.kwargs.get('mtp_num_layers') + + +def test_get_mcore_model_config_prefers_n_routed_experts(monkeypatch): + _patch_model_config(monkeypatch) + hf_config = PretrainedConfig(model_type="glm_moe_dsa", num_experts=256, n_routed_experts=16) + + config = utils.get_mcore_model_config(_make_args(), hf_config) + + assert config.kwargs['num_moe_experts'] == 16 + + +def test_get_mcore_model_config_keeps_explicit_mtp_num_layers(monkeypatch): + _patch_model_config(monkeypatch) + hf_config = PretrainedConfig(num_nextn_predict_layers=1) + + for depth in (0, 1, 2): + config = utils.get_mcore_model_config(_make_args(mtp_num_layers=depth), hf_config) + assert config.kwargs['mtp_num_layers'] == depth + + +def test_get_padding_to_sequence_parallel_uses_tp_times_two(): + args = SimpleNamespace( + tensor_model_parallel_size=2, + sequence_parallel=True, + context_parallel_size=1, + fp8_recipe='delayed', + fp8_format=None, + fp8=None, + fp4_format=None, + fp4=None, + attention_backend='unfused', + ) + assert get_padding_to(args) == 2 + args.megatron_extra_kwargs = {"dsa_accuracy_compatible": True} + assert get_padding_to(args) == 4 + args.megatron_extra_kwargs = {"dsa_accuracy_compatible": False} + assert get_padding_to(args) == 2 + seq_len = 57 + assert math.ceil(seq_len / 4) * 4 == 60 + assert math.ceil(seq_len / 2) * 2 == 58 + + +def test_dsa_index_share_allows_recompute_none(): + config = SimpleNamespace( + experimental_attention_variant='dsa', + dsa_indexer_topk_freq=4, + recompute_granularity='none', + ) + + utils._check_dsa_index_share_recompute(config) + + +def test_dsa_backend_forced_to_local_spec_when_accuracy_compatible(monkeypatch): + from megatron.core.models.backends import LocalSpecProvider + from megatron.core.models.gpt import experimental_attention_variant_module_specs as eav + + import swift.megatron.init as init + + monkeypatch.setattr(init, '_use_accuracy_compatible_enabled', lambda: True) + init._patch_mcore_bridge_disable_te() + provider = eav._get_backend_spec_provider(SimpleNamespace(dsa_accuracy_compatible=True)) + assert isinstance(provider, LocalSpecProvider) + assert hasattr(provider, 'linear') + assert provider.linear() is not provider.column_parallel_linear() + + +def test_dsa_index_share_rejects_selective_recompute(): + config = SimpleNamespace( + experimental_attention_variant='dsa', + dsa_indexer_topk_freq=4, + recompute_granularity='selective', + ) + + try: + utils._check_dsa_index_share_recompute(config) + except ValueError as error: + assert 'Set recompute_granularity=none' in str(error) + else: + raise AssertionError('expected DSA index sharing with selective recompute to fail closed') + + +def test_local_dense_norm_binding_keeps_other_parameter_keys(): + tree = ast.parse(inspect.getsource(_patch_mcore_bridge_disable_te)) + node = next(node for node in ast.walk(tree) if isinstance(node, ast.FunctionDef) and node.name == '_set_state_dict') + namespace = {'origin_set_state_dict': lambda *args, **kwargs: (args, kwargs)} + exec(compile(ast.Module(body=[node], type_ignores=[]), '', 'exec'), namespace) + bind = namespace['_set_state_dict'] + local = SimpleNamespace(mlp=SimpleNamespace(linear_fc1=SimpleNamespace())) + fused = SimpleNamespace(mlp=SimpleNamespace(linear_fc1=SimpleNamespace(layer_norm_weight=object()))) + for layer, key, expected in ( + (local, 'mlp.linear_fc1.layer_norm_weight', 'pre_mlp_layernorm.weight'), + (None, 'mlp.linear_fc1.layer_norm_weight', 'pre_mlp_layernorm.weight'), + (fused, 'mlp.linear_fc1.layer_norm_weight', 'mlp.linear_fc1.layer_norm_weight'), + (local, 'mlp.linear_fc1.weight', 'mlp.linear_fc1.weight'), + ): + args, kwargs = bind(object(), layer, key, {}, 'post_attention_layernorm.weight', True, offset=0.5) + assert args[2] == expected + assert args[4:] == ('post_attention_layernorm.weight', True) + assert kwargs == {'offset': 0.5} diff --git a/tests/megatron/test_model_repro_machine_outputs.py b/tests/megatron/test_model_repro_machine_outputs.py new file mode 100644 index 0000000000..75fccb4410 --- /dev/null +++ b/tests/megatron/test_model_repro_machine_outputs.py @@ -0,0 +1,74 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Production machine-output and checkpoint-path contracts without GPU imports.""" +import ast +import hashlib +import json +import os +import tempfile +import unittest +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch + +ROOT = Path(__file__).resolve().parents[2] +PRINT = ROOT / 'swift/megatron/callbacks/print.py' +BASE = ROOT / 'swift/megatron/trainers/base.py' +namespace = {'os': os, 'Path': Path, 'hashlib': hashlib, 'json': json} +functions = [ + n for n in ast.parse(PRINT.read_text()).body + if isinstance(n, ast.FunctionDef) and n.name in ('raw_loss_event', 'machine_loss_payload', '_write_json') +] +exec(compile(ast.Module(body=functions, type_ignores=[]), str(PRINT), 'exec'), namespace) +trainer = next( + n for n in ast.parse(BASE.read_text()).body if isinstance(n, ast.ClassDef) and n.name == 'BaseMegatronTrainer') +save = next(n for n in trainer.body if isinstance(n, ast.FunctionDef) and n.name == 'save_checkpoint') +namespace.update(gc_collect=lambda: None, save_mcore_checkpoint=lambda *args, **kwargs: None, is_master=lambda: False) +exec(compile(ast.Module(body=[save], type_ignores=[]), str(BASE), 'exec'), namespace) + + +class MachineOutputTests(unittest.TestCase): + + def test_unrounded_training_losses_exclude_evaluation(self): + event = namespace['raw_loss_event'](1, {'loss': 0.123456789123, 'mtp_1_loss': 0.5, 'eval_loss': 10}) + self.assertEqual(event, {'step': 1, 'loss': 0.123456789123, 'mtp_1_loss': 0.5}) + self.assertIsNone(namespace['raw_loss_event'](2, {'eval_loss': 10})) + result = namespace['machine_loss_payload']([event]) + self.assertEqual(result['losses'], [0.123456789123]) + self.assertEqual(result['steps'], [1]) + self.assertNotIn('owning_cli_exit_code', result) + + def test_final_override_preserves_native_export_and_args_source(self): + for iteration, override in [(99, True), (100, True), (100, False)]: + with self.subTest(iteration=iteration, override=override), tempfile.TemporaryDirectory() as directory: + output = str(Path(directory) / 'training') + final = str(Path(directory) / 'canonical' / 'checkpoint') + args = SimpleNamespace( + output_dir=output, + train_iters=100, + save_safetensors=True, + no_save_optim=True, + tuner_type='full', + merge_lora=False) + state = SimpleNamespace(iteration=iteration, consumed_train_samples=iteration, best_global_step=None) + copies, exports = [], [] + obj = SimpleNamespace( + args=args, + state=state, + optimizer=None, + opt_param_scheduler=None, + unwrapped_models=['native-model'], + template=SimpleNamespace(processor='processor'), + bridge=SimpleNamespace(save_weights=lambda *a, **kw: exports.append((a, kw))), + copy_path=lambda source, target: copies.append((source, target))) + environment = {'MODEL_REPRO_CHECKPOINT_DIR': final} if override else {} + with patch.dict(os.environ, environment, clear=True): + namespace['save_checkpoint'](obj) + expected = final if override and iteration == 100 else str(Path(output) / f'checkpoint-{iteration}') + self.assertEqual(state.last_model_checkpoint, expected) + self.assertEqual(copies[0], (str(Path(output) / 'args.json'), str(Path(expected) / 'args.json'))) + self.assertEqual(exports[0][0], (['native-model'], expected)) + self.assertEqual(exports[0][1]['processor'], 'processor') + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/megatron/test_raw_loss_observability.py b/tests/megatron/test_raw_loss_observability.py new file mode 100644 index 0000000000..f7fd6b3485 --- /dev/null +++ b/tests/megatron/test_raw_loss_observability.py @@ -0,0 +1,39 @@ +from swift.megatron.callbacks.print import raw_loss_event +from swift.megatron.trainers.trainer import project_owning_loader_semantics + + +def test_raw_loss_event_preserves_unrounded_values_and_step(): + logs = { + 'loss': 12.410510059999999, + 'mtp_0_loss': 13.367947579999999, + 'eval_loss': 99.0, + } + assert raw_loss_event(1, logs) == { + 'step': 1, + 'loss': 12.410510059999999, + 'mtp_0_loss': 13.367947579999999, + } + + +def test_raw_loss_event_omits_non_loss_metrics(): + assert raw_loss_event(3, {'grad_norm': 1.0, 'learning_rate': 1e-6}) is None + + +def test_owning_loader_projection_removes_sp_padding_and_reverses_label_roll(): + input_values = list(range(57)) + [154820] + original_labels = [-100] * 13 + list(range(44)) + [-100] + shifted_labels = original_labels[1:] + original_labels[:1] + semantic_input, semantic_labels, semantic_mask = project_owning_loader_semantics( + input_values, shifted_labels, semantic_length=57, labels_were_shifted=True) + assert semantic_input == list(range(57)) + assert semantic_labels == original_labels[:57] + assert sum(semantic_mask) == 44 + + +def test_owning_loader_projection_rejects_invalid_semantic_length(): + try: + project_owning_loader_semantics([1, 2], [-100, 2], semantic_length=3) + except ValueError as exc: + assert 'invalid owning-loader semantic length' in str(exc) + else: + raise AssertionError('invalid semantic length was accepted')