diff --git a/.bandit b/.bandit index b81428751..26ecc1d22 100644 --- a/.bandit +++ b/.bandit @@ -1,3 +1,3 @@ [bandit] -exclude = tests +exclude = tests,src/axolotl/core/trainers/prm skips = B101,B615,B102,B110 diff --git a/.mypy.ini b/.mypy.ini index 558f7600a..a55f230a2 100644 --- a/.mypy.ini +++ b/.mypy.ini @@ -8,6 +8,9 @@ ignore_missing_imports = True [mypy-axolotl.monkeypatch.*] ignore_errors = True +[mypy-axolotl.core.trainers.prm.*] +ignore_errors = True + # vendored upstream modeling code [mypy-axolotl.model_support.kimi_linear.modeling_kimi] ignore_errors = True diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index a3486461c..05a9729a9 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,6 +1,8 @@ default_language_version: python: python3 +exclude: ^src/axolotl/core/trainers/prm/ + repos: - repo: https://github.com/pre-commit/pre-commit-hooks rev: v6.0.0 diff --git a/MANIFEST.in b/MANIFEST.in index 1e51d07d6..b187d4ca7 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -1,5 +1,7 @@ include README.md include LICENSE +include src/axolotl/core/trainers/prm/LICENSE +include src/axolotl/core/trainers/prm/NOTICE include VERSION include src/axolotl/utils/chat_templates/templates/*.jinja include src/axolotl/integrations/kernels/libs/scattermoe_lora/metadata.json diff --git a/_quarto.yml b/_quarto.yml index 07ef7bc79..5eea78db6 100644 --- a/_quarto.yml +++ b/_quarto.yml @@ -224,7 +224,6 @@ quartodoc: - monkeypatch.trainer.lr - monkeypatch.trainer.trl - monkeypatch.trainer.trl_vllm - - monkeypatch.trainer.utils - monkeypatch.transformers.trainer_loss_calc - monkeypatch.xformers_ - title: Utils diff --git a/docs/agents/reward_modelling.md b/docs/agents/reward_modelling.md index 468138b74..7cb73d33b 100644 --- a/docs/agents/reward_modelling.md +++ b/docs/agents/reward_modelling.md @@ -25,7 +25,7 @@ Dataset format: `{"system": "...", "input": "...", "chosen": "...", "rejected": Train a token classifier to score each reasoning step. Uses `AutoModelForTokenClassification`. -Requires `trl<=1.13.0` (`pip install trl==1.13.0`): TRL removed `trl.experimental.prm` in later versions. ORM is unaffected. +Uses the vendored Hugging Face TRL v1.13.0 PRM trainer in `src/axolotl/core/trainers/prm/` (Apache 2.0). No TRL downgrade is required. ```yaml base_model: Qwen/Qwen2.5-3B diff --git a/docs/reward_modelling.qmd b/docs/reward_modelling.qmd index 624b60e6b..17820b4a2 100644 --- a/docs/reward_modelling.qmd +++ b/docs/reward_modelling.qmd @@ -46,13 +46,7 @@ Bradley-Terry chat templates expect single-turn conversations in the following f Check out our [PRM blog](https://axolotlai.substack.com/p/process-reward-models). ::: -::: {.callout-important} -PRM requires `trl<=1.13.0`. TRL [removed](https://github.com/huggingface/trl/pull/7133) `trl.experimental.prm` in later versions. - -```bash -uv pip install trl==1.13.0 -``` -::: +Axolotl vendors the PRM trainer from [Hugging Face TRL v1.13.0](https://github.com/huggingface/trl/tree/v1.13.0/trl/experimental/prm) under Apache 2.0. PRM works with Axolotl’s pinned TRL version; no downgrade is required. Process reward models are trained using data which contains preference annotations for each step in a series of interactions. Typically, PRMs are trained to provide reward signals over each step of a reasoning trace and are used for downstream reinforcement learning. ```yaml diff --git a/pyproject.toml b/pyproject.toml index e7152d20d..68db66ef6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,16 +14,16 @@ dependencies = [ # Core ML stack "torch>=2.13.0,<=2.14.0", "packaging==26.0", - "huggingface_hub==1.18.0", - "peft==0.21.0", + "huggingface_hub>=1.31.0", + "peft==0.21.1", "tokenizers==0.23.1", - "transformers==5.17.0", + "transformers==5.18.0", "accelerate==1.15.0", - "datasets==4.8.4", - "trl==1.13.0", - "hf_xet==1.4.3", - "kernels>=0.16.0,<0.17.0", - "trackio==0.19.0", + "datasets==5.0.1", + "trl==1.14.1", + "hf_xet>=1.5.2", + "kernels>=0.17.0,<0.18.0", + "trackio==0.40.0", "typing-extensions==4.15.0", "optimum==1.16.2", "hf_transfer", @@ -65,7 +65,7 @@ dependencies = [ "antlr4-python3-runtime==4.13.2", "schedulefree==1.4.1", - "openenv-core==0.1.0", + "openenv-core>=0.1.0", # Axolotl contribs "axolotl-contribs-lgpl==0.0.7", @@ -74,13 +74,13 @@ dependencies = [ # Telemetry "posthog==6.7.11", - "mistral-common==1.11.5", + "mistral-common==1.12.0", # Platform-specific (Linux only) "bitsandbytes==0.50.2 ; sys_platform != 'darwin'", "triton>=3.4.0 ; sys_platform != 'darwin'", "xformers>=0.0.33.post2 ; sys_platform != 'darwin' and platform_machine != 'aarch64'", - "liger-kernel==0.8.1 ; sys_platform != 'darwin'", + "liger-kernel==0.8.4 ; sys_platform != 'darwin'", "torchao==0.18.0 ; sys_platform != 'darwin' and platform_machine != 'aarch64'", # Architecture-specific @@ -185,6 +185,8 @@ include-package-data = true axolotl = [ "**/*.yaml", "**/*.jinja", + "core/trainers/prm/LICENSE", + "core/trainers/prm/NOTICE", "integrations/kernels/libs/scattermoe_lora/metadata.json", ] @@ -199,6 +201,8 @@ version = { file = "VERSION" } [tool.ruff] line-length = 88 target-version = "py312" +extend-exclude = ["src/axolotl/core/trainers/prm"] +force-exclude = true [tool.ruff.lint] select = ["E", "F", "W", "C90", "B", "I"] diff --git a/src/axolotl/core/trainers/mixins/distributed_parallel.py b/src/axolotl/core/trainers/mixins/distributed_parallel.py index ad2d436ba..26b7ea799 100644 --- a/src/axolotl/core/trainers/mixins/distributed_parallel.py +++ b/src/axolotl/core/trainers/mixins/distributed_parallel.py @@ -2,6 +2,8 @@ Mixin for correctly saving fsdp """ +import os + from accelerate import PartialState from transformers import Trainer @@ -80,13 +82,27 @@ class DistributedParallelMixin(Trainer): model._axolotl_native_nvfp4_deepspeed_prepared = True return super()._wrap_model(model, *args, **kwargs) - def _clip_grad_norm(self, model): + def _expert_parallel_enabled(self) -> bool: + parallelism = getattr(self.accelerator, "parallelism_config", None) + if parallelism is not None: + return bool(getattr(parallelism, "ep_enabled", False)) + # pure EP builds no ParallelismConfig, only the env var + return int(os.environ.get("PARALLELISM_CONFIG_EP_SIZE", "1") or 1) > 1 + + def _global_mesh(self): + return getattr( + self.accelerator, + "torch_device_mesh", + getattr(getattr(self.accelerator, "state", None), "device_mesh", None), + ) + + def _cpu_offloaded(self) -> bool: from torch.distributed.fsdp import CPUOffloadPolicy plugin = getattr(getattr(self.accelerator, "state", None), "fsdp_plugin", None) - if not isinstance(getattr(plugin, "cpu_offload", None), CPUOffloadPolicy): - return super()._clip_grad_norm(model) + return isinstance(getattr(plugin, "cpu_offload", None), CPUOffloadPolicy) + def _clip_grad_norm(self, model): from axolotl.utils.gradient_clipping import ( clip_grad_norm_ep_local_shards_, clip_grad_norm_local_shards_, @@ -95,34 +111,23 @@ class DistributedParallelMixin(Trainer): ) parameters = list(model.parameters()) - if not has_cpu_offloaded_dtensor_gradients(parameters): - return super()._clip_grad_norm(model) - self.accelerator.unscale_gradients() - parallelism = getattr(self.accelerator, "parallelism_config", None) - if getattr(parallelism, "ep_enabled", False): + if self._expert_parallel_enabled(): + # experts live on their own mesh: one ownership-filtered global norm + self.accelerator.unscale_gradients() return clip_grad_norm_ep_local_shards_( parameters, self.args.max_grad_norm, ep_local_parameters=ep_local_parameter_ids(model), - global_mesh=getattr( - self.accelerator, - "torch_device_mesh", - getattr( - getattr(self.accelerator, "state", None), "device_mesh", None - ), - ), + global_mesh=self._global_mesh(), ) - return clip_grad_norm_local_shards_(parameters, self.args.max_grad_norm) + if self._cpu_offloaded() and has_cpu_offloaded_dtensor_gradients(parameters): + self.accelerator.unscale_gradients() + return clip_grad_norm_local_shards_(parameters, self.args.max_grad_norm) + return super()._clip_grad_norm(model) def _get_grad_norm(self, model, grad_norm=None): if grad_norm is not None: return super()._get_grad_norm(model, grad_norm) - from torch.distributed.fsdp import CPUOffloadPolicy - - plugin = getattr(getattr(self.accelerator, "state", None), "fsdp_plugin", None) - parameters = list(model.parameters()) - if not isinstance(getattr(plugin, "cpu_offload", None), CPUOffloadPolicy): - return super()._get_grad_norm(model, grad_norm) from axolotl.utils.gradient_clipping import ( ep_local_parameter_ids, get_grad_norm_ep_local_shards_, @@ -130,23 +135,18 @@ class DistributedParallelMixin(Trainer): has_cpu_offloaded_dtensor_parameters, ) - if not has_cpu_offloaded_dtensor_parameters(parameters): - return super()._get_grad_norm(model, grad_norm) - self.accelerator.unscale_gradients() - parallelism = getattr(self.accelerator, "parallelism_config", None) - if getattr(parallelism, "ep_enabled", False): + parameters = list(model.parameters()) + if self._expert_parallel_enabled(): + self.accelerator.unscale_gradients() return get_grad_norm_ep_local_shards_( parameters, ep_local_parameters=ep_local_parameter_ids(model), - global_mesh=getattr( - self.accelerator, - "torch_device_mesh", - getattr( - getattr(self.accelerator, "state", None), "device_mesh", None - ), - ), + global_mesh=self._global_mesh(), ) - return get_grad_norm_local_shards_(parameters) + if self._cpu_offloaded() and has_cpu_offloaded_dtensor_parameters(parameters): + self.accelerator.unscale_gradients() + return get_grad_norm_local_shards_(parameters) + return super()._get_grad_norm(model, grad_norm) def _save_model_native( self, output_dir: str | None = None, _internal_call: bool = False diff --git a/src/axolotl/core/trainers/prm/LICENSE b/src/axolotl/core/trainers/prm/LICENSE new file mode 100644 index 000000000..f577b7741 --- /dev/null +++ b/src/axolotl/core/trainers/prm/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright 2020-2026 The HuggingFace Team + + 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://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. diff --git a/src/axolotl/core/trainers/prm/NOTICE b/src/axolotl/core/trainers/prm/NOTICE new file mode 100644 index 000000000..eb87ecc6a --- /dev/null +++ b/src/axolotl/core/trainers/prm/NOTICE @@ -0,0 +1,8 @@ +PRM trainer and configuration vendored from Hugging Face TRL v1.13.0. +Source: https://github.com/huggingface/trl/tree/v1.13.0/trl/experimental/prm +Revision: 3d9261f1fec9f9a8140099c78a65c7da73dce79c +Copyright 2020-2026 The HuggingFace Team. All rights reserved. +Licensed under the Apache License, Version 2.0; see LICENSE. + +Axolotl modifications: import shared helpers from the installed TRL package, +import PRM modules locally, and apply repository formatting. diff --git a/src/axolotl/core/trainers/prm/__init__.py b/src/axolotl/core/trainers/prm/__init__.py new file mode 100644 index 000000000..01f919ebc --- /dev/null +++ b/src/axolotl/core/trainers/prm/__init__.py @@ -0,0 +1,5 @@ +"""PRM components vendored from Hugging Face TRL; see NOTICE and LICENSE.""" + +from .prm_config import PRMConfig + +__all__ = ["PRMConfig"] diff --git a/src/axolotl/core/trainers/prm/prm_config.py b/src/axolotl/core/trainers/prm/prm_config.py new file mode 100644 index 000000000..5e60969b8 --- /dev/null +++ b/src/axolotl/core/trainers/prm/prm_config.py @@ -0,0 +1,95 @@ +# Copyright 2020-2026 The HuggingFace Team. 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://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. + +# Modified by Axolotl: use installed TRL shared helpers and local PRM modules. + +from dataclasses import dataclass, field + +from trl.trainer.base_config import _BaseConfig + + +@dataclass +class PRMConfig(_BaseConfig): + r""" + Configuration class for the [`experimental.prm.PRMTrainer`]. + + This class includes only the parameters that are specific to PRM training. For a full list of training arguments, + please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this class may + differ from those in [`~transformers.TrainingArguments`]. + + Using [`~transformers.HfArgumentParser`] we can turn this class into + [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the + command line. + + Parameters: + max_length (`int` or `None`, *optional*, defaults to `1024`): + Maximum length of the sequences (prompt + completion) used for truncation. + max_completion_length (`int`, *optional*): + Maximum length of the completion used for truncation. The completion is the concatenation of the steps. + disable_dropout (`bool`, *optional*, defaults to `True`): + Whether to disable dropout in the model. + step_separator (`str`, *optional*, defaults to `"\n"`): + Separator used to separate each step of the reasoning process. + train_on_last_step_only (`bool`, *optional*, defaults to `False`): + Whether to train only on the last step. + dataset_num_proc (`int`, *optional*): + Number of processes to use for processing the dataset. + + > [!NOTE] + > These parameters have default values different from [`~transformers.TrainingArguments`]: + > - `logging_steps`: Defaults to `10` instead of `500`. + > - `gradient_checkpointing`: Defaults to `True` instead of `False`. + > - `bf16`: Defaults to `True` if `fp16` is not set, instead of `False`. + > - `learning_rate`: Defaults to `1e-5` instead of `5e-5`. + """ + + # Parameters whose default values are overridden from TrainingArguments + learning_rate: float = field( + default=1e-5, + metadata={"help": "The initial learning rate for AdamW."}, + ) + + max_length: int | None = field( + default=1024, + metadata={ + "help": "Maximum length of the sequences (prompt + completion) used for truncation." + }, + ) + max_completion_length: int | None = field( + default=None, + metadata={ + "help": "Maximum length of the completion used for truncation. The completion is the concatenation of the " + "steps." + }, + ) + disable_dropout: bool = field( + default=True, + metadata={ + "help": "Whether to disable dropout in the model and reference model." + }, + ) + step_separator: str = field( + default="\n", + metadata={ + "help": "Separator used to separate each step of the reasoning process." + }, + ) + train_on_last_step_only: bool = field( + default=False, + metadata={"help": "Whether to train only on the last step."}, + ) + dataset_num_proc: int | None = field( + default=None, + metadata={"help": "Number of processes to use for processing the dataset."}, + ) diff --git a/src/axolotl/core/trainers/prm/prm_trainer.py b/src/axolotl/core/trainers/prm/prm_trainer.py new file mode 100644 index 000000000..e073568fb --- /dev/null +++ b/src/axolotl/core/trainers/prm/prm_trainer.py @@ -0,0 +1,386 @@ +# Copyright 2020-2026 The HuggingFace Team. 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://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. + +# Modified by Axolotl: use installed TRL shared helpers and local PRM modules. + +import textwrap +from collections.abc import Callable +from itertools import chain +from pathlib import Path + +import numpy as np +import torch +import torch.nn as nn +import transformers +from accelerate import PartialState +from accelerate.logging import get_logger +from accelerate.utils import is_peft_model +from datasets import Dataset, features +from packaging.version import Version +from transformers import ( + BaseImageProcessor, + DataCollator, + DataCollatorForTokenClassification, + FeatureExtractionMixin, + PreTrainedModel, + PreTrainedTokenizerBase, + ProcessorMixin, + TrainerCallback, +) +from transformers.trainer_utils import EvalPrediction +from transformers.utils import is_peft_available +from trl.experimental.utils import prepare_peft_model +from trl.trainer.base_trainer import _BaseTrainer +from trl.trainer.utils import disable_dropout_in_model + +from .prm_config import PRMConfig + +if is_peft_available(): + from peft import PeftConfig + + +logger = get_logger(__name__) + + +def compute_accuracy(eval_pred: EvalPrediction) -> dict[str, float]: + predictions, labels = eval_pred + if predictions.ndim == 3: + # Token classification task. Shapes are (batch_size, seq_len, num_labels) and (batch_size, seq_len) + # Used to compute the accuracy in the prm_trainer. + predictions = np.argmax(predictions, axis=2) + + # Flatten the predictions and labels to remove the ignored tokens. + predictions = np.array( + [ + p + for prediction, label in zip(predictions, labels, strict=True) + for (p, lbl) in zip(prediction, label, strict=True) + if lbl != -100 + ] + ) + labels = np.array([lbl for label in labels for lbl in label if lbl != -100]) + + else: + # Here, predictions is rewards_chosen and rewards_rejected. Shapes are (batch_size, 2) and (batch_size,) + # We want to see how much of the time rewards_chosen > rewards_rejected. + equal_mask = predictions[:, 0] == predictions[:, 1] + equal_predictions_count = int(equal_mask.sum()) + + if equal_predictions_count > 0: + # Before using the logger, the accelerate state must be initialized. It'susually the case when using this + # function inside a Trainer, but it may not be the case otherwise, in particular when unit testing. + PartialState() + + logger.warning( + f"There are {equal_predictions_count} out of {len(predictions[:, 0])} instances where the predictions " + "for both options are equal. These instances are ignored in the accuracy computation.", + ) + + # Filter out equal predictions + predictions = predictions[~equal_mask] + labels = labels[~equal_mask] + + # Use the remaining predictions for accuracy calculation + predictions = np.argmax(predictions, axis=1) + + accuracy = np.array(predictions == labels, dtype=float).mean().item() + return {"accuracy": accuracy} + + +class PRMTrainer(_BaseTrainer): + """ + Initialize PRMTrainer. + + Args: + model ([`~transformers.PreTrainedModel`]): + The model to train, preferably an `AutoModelForTokenClassification`. + args ([`experimental.prm.PRMConfig`]): + The arguments to use for training. + data_collator ([`~transformers.DataCollator`]): + The data collator to use for training. If None is specified, the default data collator + ([`~transformers.DataCollatorForTokenClassification`]) will be used which will pad the sequences to the + maximum length of the sequences in the batch, given a dataset of paired sequences. + train_dataset ([`~datasets.Dataset`]): + The dataset to use for training. + eval_dataset ([`~datasets.Dataset`]): + The dataset to use for evaluation. + processing_class ([`~transformers.PreTrainedTokenizerBase`], [`~transformers.BaseImageProcessor`], [`~transformers.FeatureExtractionMixin`] or [`~transformers.ProcessorMixin`], *optional*): + Processing class used to process the data. If provided, will be used to automatically process the inputs + for the model, and it will be saved along the model to make it easier to rerun an interrupted training or + reuse the fine-tuned model. + model_init (`Callable[[], transformers.PreTrainedModel]`): + The model initializer to use for training. If None is specified, the default model initializer will be + used. + compute_metrics (`Callable[[transformers.EvalPrediction], dict]`, *optional* defaults to `compute_accuracy`): + The metrics to use for evaluation. If no metrics are specified, the default metric (`compute_accuracy`) + will be used. + callbacks (`list[transformers.TrainerCallback]`): + The callbacks to use for training. + optimizers (`tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`): + The optimizer and scheduler to use for training. + preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`): + The function to use to preprocess the logits before computing the metrics. + peft_config ([`~peft.PeftConfig`], *optional*): + The PEFT configuration to use for training. If you pass a PEFT configuration, the model will be wrapped in + a PEFT model. + """ + + _tag_names = ["trl", "prm"] + _name = "PRM" + _paper = { + "title": "Solving math word problems with process-and outcome-based feedback", + "id": "2211.14275", + # docstyle-ignore + "citation": textwrap.dedent("""\ + @article{uesato2022solving, + title = {{Solving Math Word Problems With Process- and Outcome-Based Feedback}}, + author = {Uesato, Jonathan and Kushman, Nate and Kumar, Ramana and Song, Francis and Siegel, Noah and Wang, Lisa and Creswell, Antonia and Irving, Geoffrey and Higgins, Irina}, + year = 2022, + journal = {arXiv preprint arXiv:2211.14275} + }"""), + } + + def __init__( + self, + model: PreTrainedModel | nn.Module | None = None, + args: PRMConfig | None = None, + data_collator: DataCollator | None = None, + train_dataset: Dataset | None = None, + eval_dataset: Dataset | dict[str, Dataset] | None = None, + processing_class: PreTrainedTokenizerBase + | BaseImageProcessor + | FeatureExtractionMixin + | ProcessorMixin + | None = None, + model_init: Callable[[], PreTrainedModel] | None = None, + compute_metrics: Callable[[EvalPrediction], dict] | None = None, + callbacks: list[TrainerCallback] | None = None, + optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = ( + None, + None, + ), + preprocess_logits_for_metrics: Callable[ + [torch.Tensor, torch.Tensor], torch.Tensor + ] + | None = None, + peft_config: "PeftConfig | None" = None, + ): + if train_dataset is None: + raise ValueError("`train_dataset` is required") + + # PEFT + if peft_config is not None: + if not is_peft_available(): + raise ImportError( + "You passed `peft_config` but the `peft` library is not installed. " + "Install it with `pip install trl[peft]`." + ) + if not isinstance(peft_config, PeftConfig): + raise TypeError( + f"`peft_config` must be a `peft.PeftConfig` instance (e.g. `peft.LoraConfig`), " + f"got {type(peft_config).__name__}." + ) + if peft_config is not None or is_peft_model(model): + model = prepare_peft_model(model, peft_config, args) + + # Disable dropout in the model + if args.disable_dropout: + disable_dropout_in_model(model) + + if compute_metrics is None: + compute_metrics = compute_accuracy + + if data_collator is None: + if processing_class is None: + raise ValueError( + "A processing_class must be specified when using the default DataCollatorForTokenClassification" + ) + data_collator = DataCollatorForTokenClassification(processing_class) + + if "input_ids" not in train_dataset.column_names: + with PartialState().main_process_first(): + fn_kwargs = { + "tokenizer": processing_class, + "step_separator": args.step_separator, + "max_length": args.max_length, + "max_completion_length": args.max_completion_length, + "train_on_last_step_only": args.train_on_last_step_only, + } + train_fn_kwargs = {**fn_kwargs, "is_eval": False} + train_dataset = train_dataset.map( + self.tokenize_row, + fn_kwargs=train_fn_kwargs, + num_proc=args.dataset_num_proc, + remove_columns=train_dataset.features, + desc="Tokenizing train dataset", + features=features.Features( # needed to avoid map to cast labels to bool + { + "labels": features.Sequence(features.Value("int64")), + "input_ids": features.Sequence(features.Value("int64")), + } + ), + ) + + eval_fn_kwargs = {**fn_kwargs, "is_eval": True} + if eval_dataset is not None: + eval_dataset = eval_dataset.map( + self.tokenize_row, + fn_kwargs=eval_fn_kwargs, + num_proc=args.dataset_num_proc, + remove_columns=eval_dataset.features, + desc="Tokenizing eval dataset", + features=features.Features( # needed to avoid map to cast labels to bool + { + "labels": features.Sequence(features.Value("int64")), + "input_ids": features.Sequence(features.Value("int64")), + } + ), + ) + + # Transformers explicitly set use_reentrant=True in the past to silence a PyTorch warning, but the default was + # never updated once PyTorch switched to recommending use_reentrant=False. Until that change lands upstream + # (see https://github.com/huggingface/transformers/pull/43203) and is released (most likely in 5.0.0), we + # default to the recommended non-reentrant behavior here, while preserving any user-provided value. + if args.gradient_checkpointing and Version(transformers.__version__) < Version( + "5.0.0" + ): + args.gradient_checkpointing_kwargs = ( + args.gradient_checkpointing_kwargs or {} + ) + args.gradient_checkpointing_kwargs.setdefault("use_reentrant", False) + + super().__init__( + model=model, + args=args, + data_collator=data_collator, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + model_init=model_init, + compute_metrics=compute_metrics, + callbacks=callbacks, + optimizers=optimizers, + preprocess_logits_for_metrics=preprocess_logits_for_metrics, + ) + + # Add tags for models that have been loaded with the correct transformers version + if hasattr(self.model, "add_model_tags"): + self.model.add_model_tags(self._tag_names) + + @staticmethod + def tokenize_row( + features, + tokenizer, + step_separator, + max_length, + max_completion_length, + train_on_last_step_only, + is_eval, + ): + r""" + Tokenize a row of the dataset. + + Args: + features (`dict[str, str]`): + Row of the dataset, should contain the keys `"prompt"`, `"completions"`, and `"labels"`. + tokenizer ([`~transformers.PreTrainedTokenizerBase`]): + Tokenizer used to process the data. + step_separator (`str`): + Separator between steps in the completion. + max_length (`int` or `None`): + Maximum length of the sequences (prompt + completion). If `None`, the sequences are not truncated. + max_completion_length (`int` or `None`): + Maximum length of the completion sequences. If `None`, the completion sequences are not truncated. + train_on_last_step_only (`bool`): + Whether to train only on the last step. If `True`, the labels are `-100` for all tokens except the last + token of the completion. + is_eval (`bool`): + Whether the function is used to tokenize samples from a training or an evaluation dataset. Used only if + `train_on_last_step_only` is set to `True`. + + Returns: + `dict[str, list[int]]`: + Tokenized sequences with the keys `"input_ids"`, and `"labels". + + Example: + ```python + >>> from transformers import AutoTokenizer + + >>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B") + >>> features = { + ... "prompt": "Which number is larger, 9.8 or 9.11?", + ... "completions": ["11 is greater than 8.", "Hence, 9.11 > 9.8."], + ... "labels": [True, False], + ... } + >>> PRMTrainer.tokenize_row( + ... features, tokenizer, "\n", max_completion_length=None, train_on_last_step_only=False, is_eval=False + ... ) + {'input_ids': [23085, 1372, 374, 8131, 11, 220, 24, 13, 23, 476, 220, 24, 13, 16, 16, 30, 16, 16, 374, 7046, 1091, 220, 23, 13, 198, 39, 763, 11, 220, 24, 13, 16, 16, 861, 220, 24, 13, 23, 13, 198], + 'labels': [-100, -100, -100, -100, -100, -100, -100, -100, 1, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, 0]} + ``` + """ + # Tokenize the prompt and completions + prompt_ids = tokenizer(features["prompt"], add_special_tokens=False)[ + "input_ids" + ] + completions_ids = [ + tokenizer(completion, add_special_tokens=False)["input_ids"] + for completion in features["completions"] + ] + if train_on_last_step_only and not is_eval: + labels = [-100] * (len(features["labels"]) - 1) + [ + int(features["labels"][-1]) + ] + else: + labels = [int(label) for label in features["labels"]] + + # Get the ID of the separator token and add it to the completions + separator_ids = tokenizer.encode(step_separator, add_special_tokens=False) + completions_ids = [completion + separator_ids for completion in completions_ids] + + # Create the label + labels = [ + [-100] * (len(completion) - 1) + [label] + for completion, label in zip(completions_ids, labels, strict=True) + ] + + # Join the completions and labels steps + completion_ids = list(chain(*completions_ids)) + labels = list(chain(*labels)) + + if tokenizer.bos_token_id is not None: + prompt_ids = [tokenizer.bos_token_id] + prompt_ids + + # Truncate completion sequences + if max_completion_length is not None: + completion_ids = completion_ids[:max_completion_length] + labels = labels[:max_completion_length] + + input_ids = prompt_ids + completion_ids + labels = [-100] * len(prompt_ids) + labels + + if max_length is not None: + input_ids = input_ids[:max_length] + labels = labels[:max_length] + + return {"input_ids": input_ids, "labels": labels} + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) diff --git a/src/axolotl/core/trainers/trl.py b/src/axolotl/core/trainers/trl.py index d9e4338ee..7fe334c56 100644 --- a/src/axolotl/core/trainers/trl.py +++ b/src/axolotl/core/trainers/trl.py @@ -7,19 +7,7 @@ from trl.experimental.orpo import ORPOTrainer from axolotl.core.trainers.mixins import DistributedParallelMixin, RngLoaderMixin from axolotl.core.trainers.mixins.optimizer import OptimizerInitMixin, OptimizerMixin from axolotl.core.trainers.mixins.scheduler import SchedulerMixin -from axolotl.core.training_args import PRM_TRL_ERROR - -try: - from trl.experimental.prm import PRMTrainer -except ImportError: - - class PRMTrainer: # type: ignore[no-redef] - """Placeholder so this module stays importable without trl's PRM.""" - - # __new__, not __init__: the trainer mixins put transformers' Trainer - # ahead of this placeholder in the MRO. - def __new__(cls, *args, **kwargs): - raise ImportError(PRM_TRL_ERROR) +from axolotl.core.trainers.prm.prm_trainer import PRMTrainer class AxolotlORPOTrainer( @@ -91,7 +79,7 @@ class AxolotlPRMTrainer( PRMTrainer, ): """ - Extend the base trl.PRMTrainer for axolotl helpers + Extend the vendored PRMTrainer for axolotl helpers """ tag_names = ["axolotl", "prm"] diff --git a/src/axolotl/core/training_args.py b/src/axolotl/core/training_args.py index 1b6e646ae..2e9cdc3ea 100644 --- a/src/axolotl/core/training_args.py +++ b/src/axolotl/core/training_args.py @@ -12,25 +12,9 @@ from trl import KTOConfig, RewardConfig from trl.experimental.cpo import CPOConfig from trl.experimental.orpo import ORPOConfig +from axolotl.core.trainers.prm import PRMConfig from axolotl.integrations.config import merge_training_args -PRM_TRL_ERROR = ( - "PRM training requires `trl<=1.13.0` (`pip install trl==1.13.0`); " - "`trl.experimental.prm` was removed in later versions." -) - -try: - from trl.experimental.prm import PRMConfig -except ImportError: - - @dataclass - class PRMConfig(TrainingArguments): # type: ignore[no-redef] - """Placeholder so this module stays importable without trl's PRM.""" - - def __post_init__(self): - raise ImportError(PRM_TRL_ERROR) - - AxolotlTrainingMixins: Type = merge_training_args() diff --git a/src/axolotl/integrations/expert_parallel/experts_fn.py b/src/axolotl/integrations/expert_parallel/experts_fn.py index 14f44d45d..928b651e7 100644 --- a/src/axolotl/integrations/expert_parallel/experts_fn.py +++ b/src/axolotl/integrations/expert_parallel/experts_fn.py @@ -179,7 +179,8 @@ def _normalize_sentinels(recv_topk_idx, recv_topk_weights, num_local_experts): """Map the dispatch's ``-1`` remote slots to ``num_local_experts`` with weight 0. transformers' grouped_mm / sonicmoe / deepgemm drop ids ``>= num_experts`` natively; - batched_mm clamps them and relies on the zero weight. + batched_mm clamps them and relies on the zero weight. Both need the module's + ``_is_expert_parallel`` set first, which ``_ep_forward`` does. """ remote = recv_topk_idx < 0 return ( @@ -389,6 +390,7 @@ def _ep_forward( ep_backend = _BACKENDS[backend] local_fn = resolve_local_implementation(local) _maybe_install_decorator_attrs(self) + self._is_expert_parallel = True original_dtype = None if ep_backend.cast_bf16 and hidden_states.dtype != torch.bfloat16: original_dtype = hidden_states.dtype diff --git a/src/axolotl/integrations/expert_parallel/shard.py b/src/axolotl/integrations/expert_parallel/shard.py index b1e5281b3..66c111390 100644 --- a/src/axolotl/integrations/expert_parallel/shard.py +++ b/src/axolotl/integrations/expert_parallel/shard.py @@ -227,6 +227,7 @@ def shard_expert_weights(model, ep_group) -> int: model._ep_num_experts_global = E # Override the kernel's view of num_experts for grouped_mm/scattermoe bucketing. module.num_experts = E_local + module._is_expert_parallel = True # Mark sharded params as DDP-ignored — they hold rank-specific content # and must NOT be broadcast from rank 0 at DDP construction. diff --git a/src/axolotl/loaders/patch_manager.py b/src/axolotl/loaders/patch_manager.py index bc664c412..39cd4b27a 100644 --- a/src/axolotl/loaders/patch_manager.py +++ b/src/axolotl/loaders/patch_manager.py @@ -153,7 +153,6 @@ class PatchManager: self._apply_voxtral_patches() self._apply_apertus_patches() self._apply_trl_vllm_patches() - self._apply_trl_trainer_utils_patches() def apply_post_plugin_pre_model_load_patches(self): """Apply post plugin-pre_model_load load patches based on config.""" @@ -1265,39 +1264,6 @@ class PatchManager: patch_trl_vllm() - def _apply_trl_trainer_utils_patches(self): - """Replace trl.trainer.utils.{selective_log_softmax, entropy_from_logits} with Triton kernels.""" - if not self.cfg.rl: - return - - try: - from axolotl.monkeypatch.trainer.utils import ( - entropy_from_logits, - selective_log_softmax, - ) - except (ImportError, ModuleNotFoundError): - LOG.warning("Triton not available — skipping trl.trainer.utils patches") - return - - import trl.trainer.utils - - # Guard against repeated calls: only stash the original if trl still - # points at its own implementation (not our wrapper). - if trl.trainer.utils.selective_log_softmax is not selective_log_softmax: - from axolotl.monkeypatch.trainer import utils as _axolotl_trainer_utils - - _axolotl_trainer_utils.selective_log_softmax_original = ( - trl.trainer.utils.selective_log_softmax - ) - trl.trainer.utils.selective_log_softmax = selective_log_softmax - - if trl.trainer.utils.entropy_from_logits is not entropy_from_logits: - trl.trainer.utils.entropy_from_logits = entropy_from_logits - - LOG.info( - "Patched trl.trainer.utils with Triton selective_log_softmax and entropy_from_logits" - ) - def _apply_scaling_softmax_patch(self, model: PreTrainedModel): """Apply Scaling Softmax (SSMax) patch. Ref: https://arxiv.org/abs/2501.19399""" if self.cfg.scaling_softmax: diff --git a/src/axolotl/monkeypatch/accelerate/parallelism_config.py b/src/axolotl/monkeypatch/accelerate/parallelism_config.py index faf273b98..e725cd700 100644 --- a/src/axolotl/monkeypatch/accelerate/parallelism_config.py +++ b/src/axolotl/monkeypatch/accelerate/parallelism_config.py @@ -282,80 +282,6 @@ def patch_parallelism_config(): ParallelismConfig._validate_accelerator = _validate_accelerator AcceleratorState.is_fsdp2 = property(patched_is_fsdp2) patch_prepare_data_loader_for_ep() - patch_clip_grad_norm_for_ep() - - -def _ep_aware_clip_grad_norm( - parameters, - max_norm, - norm_type=2.0, - *, - ep_local_parameters=frozenset(), - global_mesh=None, -): - """Clip EP/FSDP gradients by one ownership-filtered global norm.""" - from axolotl.utils.gradient_clipping import clip_grad_norm_ep_local_shards_ - - return clip_grad_norm_ep_local_shards_( - parameters, - max_norm, - norm_type=norm_type, - ep_local_parameters=set(ep_local_parameters), - global_mesh=global_mesh, - ) - - -def patch_clip_grad_norm_for_ep(): - """Replace `Accelerator.clip_grad_norm_` with the EP-aware version when - the active parallelism composes `ep` with FSDP2 (`dp_shard`, `cp`, or pure EP - with FSDP-sharded dense params): the FSDP+EP composition produces multi-mesh - DTensor grads — the experts shard on the dp_shard/cp subgroup (or a per-rank - mesh), the non-experts on the flattened mesh, so the stock `clip_grad_norm_` - can't stack their per-param norms together. - """ - from accelerate import Accelerator - - if getattr(Accelerator, "_AXOLOTL_EP_CLIP_PATCHED", False): - return - orig = Accelerator.clip_grad_norm_ - - def patched_clip_grad_norm_(self, parameters, max_norm, norm_type=2): - pc = getattr(self, "parallelism_config", None) - # Pure EP builds no ParallelismConfig (only the env var), yet it still mixes meshes: - # dense params on the flat world mesh, trainable experts on a per-rank mesh whose - # `full_tensor()` sees one rank's experts, so the stock clip gets a per-rank norm. - ep_enabled = ( - getattr(pc, "ep_enabled", False) - if pc is not None - else int(os.environ.get("PARALLELISM_CONFIG_EP_SIZE", "1") or 1) > 1 - ) - if ep_enabled and ( - getattr(pc, "dp_shard_enabled", False) - or getattr(pc, "cp_enabled", False) - or getattr(self, "is_fsdp2", False) - ): - self.unscale_gradients() - params = list(parameters) - from axolotl.utils.gradient_clipping import ep_local_parameter_ids - - ep_local = set() - for model in getattr(self, "_models", ()): - ep_local.update(ep_local_parameter_ids(model)) - return _ep_aware_clip_grad_norm( - params, - max_norm, - norm_type=norm_type, - ep_local_parameters=ep_local, - global_mesh=getattr( - self, - "torch_device_mesh", - getattr(getattr(self, "state", None), "device_mesh", None), - ), - ) - return orig(self, parameters, max_norm, norm_type=norm_type) - - Accelerator.clip_grad_norm_ = patched_clip_grad_norm_ - Accelerator._AXOLOTL_EP_CLIP_PATCHED = True def _patched_prepare_data_loader_factory(orig_fn): diff --git a/src/axolotl/monkeypatch/torchao_tp_lora.py b/src/axolotl/monkeypatch/torchao_tp_lora.py index 86de52eb0..d72e0cf78 100644 --- a/src/axolotl/monkeypatch/torchao_tp_lora.py +++ b/src/axolotl/monkeypatch/torchao_tp_lora.py @@ -138,6 +138,30 @@ def _merge_aware_unsupported_reason(module) -> str | None: return None +def _localize_peft_tp_factors(module, layout: str) -> None: + """PEFT >= 0.21.1 wraps the sharded factor as a DTensor and installs transformers' + TP forward on it; this path computes with local factors and does its own reductions.""" + sharded = ("lora_B", 0) if layout == "colwise" else ("lora_A", 1) + for factor_name in ("lora_A", "lora_B"): + for factor in getattr(module, factor_name).values(): + weight = factor.weight + if type(weight).__name__ == "DTensor": + placement = weight.placements[-1] + expected_dim = sharded[1] if factor_name == sharded[0] else None + if ( + expected_dim is None + or getattr(placement, "dim", None) != expected_dim + ): + raise ValueError( + f"Native NVFP4 TP {layout} LoRA expects {sharded[0]} sharded on dim " + f"{sharded[1]}, got {factor_name} with placement {placement}" + ) + factor.weight = torch.nn.Parameter( + weight.to_local(), requires_grad=weight.requires_grad + ) + factor.__dict__.pop("forward", None) + + def _validate_ordinary_lora(module, layout: str) -> None: if getattr(module, "lora_variant", None) or any( getattr(module, "use_dora", {}).values() @@ -308,6 +332,7 @@ def prepare_native_nvfp4_tp_lora( "Native NVFP4 TP LoRA only supports colwise and rowwise TP plans" ) layout = _tp_lora_layout(weight, plan) + _localize_peft_tp_factors(module, layout) _validate_ordinary_lora(module, layout) unavailable = _merge_aware_unsupported_reason(module) if merge_aware else None if unavailable is not None: diff --git a/src/axolotl/monkeypatch/trainer/__init__.py b/src/axolotl/monkeypatch/trainer/__init__.py index 27edc63c3..a4693a883 100644 --- a/src/axolotl/monkeypatch/trainer/__init__.py +++ b/src/axolotl/monkeypatch/trainer/__init__.py @@ -1,3 +1 @@ -from .utils import entropy_from_logits, selective_log_softmax - -__all__ = ["entropy_from_logits", "selective_log_softmax"] +"""Trainer-side monkeypatches.""" diff --git a/src/axolotl/monkeypatch/trainer/utils.py b/src/axolotl/monkeypatch/trainer/utils.py deleted file mode 100644 index 18a7d0cc7..000000000 --- a/src/axolotl/monkeypatch/trainer/utils.py +++ /dev/null @@ -1,510 +0,0 @@ -# Copyright 2026 Axolotl AI. 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://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 -import torch.nn.functional as F -import triton -import triton.language as tl - -from axolotl.kernels.op_registry import register_kernel_op - - -@triton.jit -def _entropy_online_kernel( - logits_ptr, - output_ptr, - stride_row, - V: tl.constexpr, - BLOCK_V: tl.constexpr, -): - """Online entropy: single pass with running max correction.""" - row = tl.program_id(0) - row_ptr = logits_ptr + tl.cast(row, tl.int64) * stride_row - - running_max = tl.full([], float("-inf"), dtype=tl.float32) - running_sum_exp = tl.full([], 0.0, dtype=tl.float32) - running_weighted = tl.full([], 0.0, dtype=tl.float32) - - for v_start in range(0, V, BLOCK_V): - offs = v_start + tl.arange(0, BLOCK_V) - mask = offs < V - x = tl.load(row_ptr + offs, mask=mask, other=float("-inf")).to(tl.float32) - - block_max = tl.max(x, axis=0) - new_max = tl.maximum(running_max, block_max) - - correction = tl.exp(running_max - new_max) - running_sum_exp = running_sum_exp * correction - running_weighted = running_weighted * correction - - exp_x = tl.exp(x - new_max) - exp_x = tl.where(mask, exp_x, 0.0) - x = tl.where(mask, x, 0.0) - running_sum_exp += tl.sum(exp_x, axis=0) - running_weighted += tl.sum(exp_x * x, axis=0) - - running_max = new_max - - entropy = tl.log(running_sum_exp) + running_max - running_weighted / running_sum_exp - tl.store(output_ptr + row, entropy) - - -@triton.jit -def _entropy_online_kernel_strided( - logits_ptr, - output_ptr, - stride_outer, - stride_inner, - n_inner, - row_offset, - V: tl.constexpr, - BLOCK_V: tl.constexpr, -): - """Online entropy for non-contiguous 3D (B, L, V) tensors.""" - local_row = tl.program_id(0) - row = local_row + row_offset - outer_idx = row // n_inner - inner_idx = row % n_inner - off = outer_idx.to(tl.int64) * stride_outer + inner_idx.to(tl.int64) * stride_inner - row_ptr = logits_ptr + off - - running_max = tl.full([], float("-inf"), dtype=tl.float32) - running_sum_exp = tl.full([], 0.0, dtype=tl.float32) - running_weighted = tl.full([], 0.0, dtype=tl.float32) - - for v_start in range(0, V, BLOCK_V): - offs = v_start + tl.arange(0, BLOCK_V) - mask = offs < V - x = tl.load(row_ptr + offs, mask=mask, other=float("-inf")).to(tl.float32) - - block_max = tl.max(x, axis=0) - new_max = tl.maximum(running_max, block_max) - - correction = tl.exp(running_max - new_max) - running_sum_exp = running_sum_exp * correction - running_weighted = running_weighted * correction - - exp_x = tl.exp(x - new_max) - exp_x = tl.where(mask, exp_x, 0.0) - x = tl.where(mask, x, 0.0) - running_sum_exp += tl.sum(exp_x, axis=0) - running_weighted += tl.sum(exp_x * x, axis=0) - - running_max = new_max - - entropy = tl.log(running_sum_exp) + running_max - running_weighted / running_sum_exp - tl.store(output_ptr + local_row, entropy) - - -@register_kernel_op("entropy_from_logits") -def _entropy_from_logits_op(logits: torch.Tensor, chunk_size: int) -> torch.Tensor: - """CUDA entropy launches; returns flat float32 `(N,)` output.""" - V = logits.shape[-1] - N = 1 - for s in logits.shape[:-1]: - N *= s - - output = torch.empty(N, device=logits.device, dtype=torch.float32) - - BLOCK_V = 4096 - MAX_GRID_CONTIG = 8192 - MAX_GRID_STRIDED = 2048 - - # Vocab (last) dim must be contiguous for coalesced loads - if logits.stride(-1) != 1: - logits = logits.contiguous() - - if logits.is_contiguous(): - flat_logits = logits.reshape(-1, V) - stride = flat_logits.stride(0) - for start in range(0, N, MAX_GRID_CONTIG): - n_rows = min(MAX_GRID_CONTIG, N - start) - _entropy_online_kernel[(n_rows,)]( - flat_logits[start], output[start], stride, V=V, BLOCK_V=BLOCK_V - ) - elif logits.ndim == 3: - stride_outer = logits.stride(0) - stride_inner = logits.stride(1) - n_inner = logits.shape[1] - for start in range(0, N, MAX_GRID_STRIDED): - n_rows = min(MAX_GRID_STRIDED, N - start) - _entropy_online_kernel_strided[(n_rows,)]( - logits, - output[start], - stride_outer, - stride_inner, - n_inner, - start, - V=V, - BLOCK_V=BLOCK_V, - ) - else: - logits = logits.contiguous() - flat_logits = logits.reshape(-1, V) - stride = flat_logits.stride(0) - for start in range(0, N, MAX_GRID_CONTIG): - n_rows = min(MAX_GRID_CONTIG, N - start) - _entropy_online_kernel[(n_rows,)]( - flat_logits[start], output[start], stride, V=V, BLOCK_V=BLOCK_V - ) - - return output - - -@_entropy_from_logits_op.register_fake -def _(logits, chunk_size): - N = 1 - for s in logits.shape[:-1]: - N *= s - return torch.empty(N, device=logits.device, dtype=torch.float32) - - -def entropy_from_logits(logits: torch.Tensor, chunk_size: int = 128) -> torch.Tensor: - """Triton-fused entropy (online single-pass). Handles non-contiguous tensors without copying.""" - original_shape = logits.shape[:-1] - - if not logits.is_cuda: - # CPU fallback: stable entropy via log_softmax - logp = F.log_softmax(logits.float(), dim=-1) - ent = -(logp.exp() * logp).sum(dim=-1) - return ent.to(logits.dtype).reshape(original_shape) - - output = _entropy_from_logits_op(logits, chunk_size) - return output.to(logits.dtype).reshape(original_shape) - - -# --------------------------------------------------------------------------- -# selective_log_softmax — fused forward + backward Triton kernels -# --------------------------------------------------------------------------- - - -def selective_log_softmax_original(logits, index) -> torch.Tensor: - """Original selective_log_softmax (reference/fallback).""" - squeeze = index.ndim == logits.ndim - 1 - if squeeze: - index = index.unsqueeze(-1) - - if logits.dtype in [torch.float32, torch.float64]: - selected_logits = torch.gather(logits, dim=-1, index=index) - logsumexp_values = torch.stack([torch.logsumexp(lg, dim=-1) for lg in logits]) - per_token_logps = selected_logits - logsumexp_values.unsqueeze(-1) - else: - per_token_logps = [] - for row_logits, row_labels in zip(logits, index, strict=True): - row_logps = F.log_softmax(row_logits, dim=-1) - row_per_token_logps = row_logps.gather(dim=-1, index=row_labels) - per_token_logps.append(row_per_token_logps) - per_token_logps = torch.stack(per_token_logps) - - if squeeze: - per_token_logps = per_token_logps.squeeze(-1) - - return per_token_logps - - -@triton.jit -def _selective_logsoftmax_fwd_kernel( - logits_ptr, - index_ptr, - output_ptr, - logsumexp_ptr, - stride_logits_row, - stride_index_row, - stride_output_row, - actual_K, - K_BLOCK: tl.constexpr, - V: tl.constexpr, - BLOCK_V: tl.constexpr, -): - """Forward: online logsumexp + gather. Saves logsumexp for backward.""" - row = tl.program_id(0) - logits_row_ptr = logits_ptr + tl.cast(row, tl.int64) * stride_logits_row - - # Online logsumexp - running_max = tl.full([], float("-inf"), dtype=tl.float32) - running_sum_exp = tl.full([], 0.0, dtype=tl.float32) - - for v_start in range(0, V, BLOCK_V): - offs = v_start + tl.arange(0, BLOCK_V) - mask = offs < V - x = tl.load(logits_row_ptr + offs, mask=mask, other=float("-inf")).to( - tl.float32 - ) - - block_max = tl.max(x, axis=0) - new_max = tl.maximum(running_max, block_max) - running_sum_exp = running_sum_exp * tl.exp(running_max - new_max) - - exp_x = tl.exp(x - new_max) - exp_x = tl.where(mask, exp_x, 0.0) - running_sum_exp += tl.sum(exp_x, axis=0) - running_max = new_max - - lse = tl.log(running_sum_exp) + running_max - tl.store(logsumexp_ptr + row, lse) - - # Gather and subtract - index_row_ptr = index_ptr + tl.cast(row, tl.int64) * stride_index_row - output_row_ptr = output_ptr + tl.cast(row, tl.int64) * stride_output_row - - k_offs = tl.arange(0, K_BLOCK) - k_mask = k_offs < actual_K - indices = tl.load(index_row_ptr + k_offs, mask=k_mask, other=0).to(tl.int64) - valid_mask = k_mask & (indices >= 0) & (indices < V) - safe_indices = tl.where(valid_mask, indices, 0) - selected = tl.load(logits_row_ptr + safe_indices, mask=valid_mask, other=0.0).to( - tl.float32 - ) - tl.store(output_row_ptr + k_offs, selected - lse, mask=valid_mask) - - -@triton.jit -def _selective_logsoftmax_bwd_kernel( - grad_output_ptr, - logits_ptr, - index_ptr, - logsumexp_ptr, - grad_logits_ptr, - stride_grad_out_row, - stride_logits_row, - stride_index_row, - stride_grad_logits_row, - actual_K, - K_BLOCK: tl.constexpr, - V: tl.constexpr, - BLOCK_V: tl.constexpr, -): - """Backward: d_logits[j] = -softmax(x)[j] * sum(grad_out) + (grad_out[k] if j == index[k]). - - Single fused pass over V. For each tile, computes the base gradient and adds - scatter contributions inline by checking which indices fall in the current tile. - No separate scatter pass — no read-after-write issues. - """ - row = tl.program_id(0) - logits_row_ptr = logits_ptr + tl.cast(row, tl.int64) * stride_logits_row - grad_logits_row_ptr = ( - grad_logits_ptr + tl.cast(row, tl.int64) * stride_grad_logits_row - ) - grad_out_row_ptr = grad_output_ptr + tl.cast(row, tl.int64) * stride_grad_out_row - index_row_ptr = index_ptr + tl.cast(row, tl.int64) * stride_index_row - - lse = tl.load(logsumexp_ptr + row).to(tl.float32) - - # Load grad_output and indices (K_BLOCK elements, masked) - k_offs = tl.arange(0, K_BLOCK) - k_mask = k_offs < actual_K - grad_out = tl.load(grad_out_row_ptr + k_offs, mask=k_mask, other=0.0).to(tl.float32) - indices = tl.load( - index_row_ptr + k_offs, mask=k_mask, other=-1 - ) # -1 = never matches - valid_mask = k_mask & (indices >= 0) & (indices < V) - grad_out = tl.where(valid_mask, grad_out, 0.0) - indices = tl.where(valid_mask, indices, -1) - grad_sum = tl.sum(grad_out, axis=0) - - # Fused pass: for each tile, compute -softmax * grad_sum + scatter - for v_start in range(0, V, BLOCK_V): - offs = v_start + tl.arange(0, BLOCK_V) # [BLOCK_V] - mask = offs < V - x = tl.load(logits_row_ptr + offs, mask=mask, other=0.0).to(tl.float32) - softmax_j = tl.exp(x - lse) - softmax_j = tl.where(mask, softmax_j, 0.0) - grad_j = -softmax_j * grad_sum - - # Scatter: check which selected indices fall in this tile - # offs: [BLOCK_V], indices: [K_BLOCK] - # Broadcast: offs[:, None] == indices[None, :] → [BLOCK_V, K_BLOCK] - match = offs[:, None] == indices[None, :] # [BLOCK_V, K_BLOCK] - # Sum grad_out contributions: for each position j, sum grad_out[k] where index[k]==j - scatter_contrib = tl.sum( - tl.where(match, grad_out[None, :], 0.0), axis=1 - ) # [BLOCK_V] - grad_j += scatter_contrib - - tl.store(grad_logits_row_ptr + offs, grad_j, mask=mask) - - -@register_kernel_op("selective_log_softmax_fwd") -def _selective_log_softmax_fwd_op( - flat_logits: torch.Tensor, - flat_index: torch.Tensor, - K: int, - K_BLOCK: int, - V: int, - BLOCK_V: int, - MAX_GRID: int, -) -> tuple[torch.Tensor, torch.Tensor]: - N = flat_logits.shape[0] - output = torch.empty(N, K_BLOCK, device=flat_logits.device, dtype=torch.float32) - logsumexp = torch.empty(N, device=flat_logits.device, dtype=torch.float32) - - for start in range(0, N, MAX_GRID): - n_rows = min(MAX_GRID, N - start) - _selective_logsoftmax_fwd_kernel[(n_rows,)]( - flat_logits[start], - flat_index[start], - output[start], - logsumexp[start], - flat_logits.stride(0), - flat_index.stride(0), - output.stride(0), - K, - K_BLOCK=K_BLOCK, - V=V, - BLOCK_V=BLOCK_V, - ) - - return output, logsumexp - - -@_selective_log_softmax_fwd_op.register_fake -def _(flat_logits, flat_index, K, K_BLOCK, V, BLOCK_V, MAX_GRID): - N = flat_logits.shape[0] - return ( - torch.empty(N, K_BLOCK, device=flat_logits.device, dtype=torch.float32), - torch.empty(N, device=flat_logits.device, dtype=torch.float32), - ) - - -@register_kernel_op("selective_log_softmax_bwd") -def _selective_log_softmax_bwd_op( - grad_output: torch.Tensor, - flat_logits: torch.Tensor, - flat_index: torch.Tensor, - logsumexp: torch.Tensor, - K: int, - K_BLOCK: int, - V: int, - BLOCK_V: int, - MAX_GRID: int, -) -> torch.Tensor: - N = flat_logits.shape[0] - - grad_logits = torch.empty_like(flat_logits) - - # grad_output may have K_BLOCK cols; backward kernel reads actual_K - grad_output_contig = grad_output.contiguous() - - for start in range(0, N, MAX_GRID): - n_rows = min(MAX_GRID, N - start) - _selective_logsoftmax_bwd_kernel[(n_rows,)]( - grad_output_contig[start], - flat_logits[start], - flat_index[start], - logsumexp[start], - grad_logits[start], - grad_output_contig.stride(0), - flat_logits.stride(0), - flat_index.stride(0), - grad_logits.stride(0), - K, - K_BLOCK=K_BLOCK, - V=V, - BLOCK_V=BLOCK_V, - ) - - return grad_logits - - -@_selective_log_softmax_bwd_op.register_fake -def _( - grad_output, flat_logits, flat_index, logsumexp, K, K_BLOCK, V, BLOCK_V, MAX_GRID -): - return torch.empty_like(flat_logits) - - -class _SelectiveLogSoftmaxTriton(torch.autograd.Function): - @staticmethod - def forward(ctx, flat_logits, flat_index, K, K_BLOCK, V, BLOCK_V, MAX_GRID): - output, logsumexp = _selective_log_softmax_fwd_op( - flat_logits, flat_index, K, K_BLOCK, V, BLOCK_V, MAX_GRID - ) - - ctx.save_for_backward(flat_logits, flat_index, logsumexp) - ctx.K = K - ctx.K_BLOCK = K_BLOCK - ctx.V = V - ctx.BLOCK_V = BLOCK_V - ctx.MAX_GRID = MAX_GRID - return output - - @staticmethod - def backward(ctx, grad_output): - flat_logits, flat_index, logsumexp = ctx.saved_tensors - K, K_BLOCK, V, BLOCK_V, MAX_GRID = ( - ctx.K, - ctx.K_BLOCK, - ctx.V, - ctx.BLOCK_V, - ctx.MAX_GRID, - ) - grad_logits = _selective_log_softmax_bwd_op( - grad_output, - flat_logits, - flat_index, - logsumexp, - K, - K_BLOCK, - V, - BLOCK_V, - MAX_GRID, - ) - - # Return grads for: flat_logits, flat_index, K, K_BLOCK, V, BLOCK_V, MAX_GRID - return grad_logits, None, None, None, None, None, None - - -def selective_log_softmax(logits, index) -> torch.Tensor: - """ - Fused selective_log_softmax with Triton forward+backward kernels. - - Equivalent to: torch.gather(logits.log_softmax(-1), dim=-1, index=index) - """ - squeeze = index.ndim == logits.ndim - 1 - if squeeze: - index = index.unsqueeze(-1) - - if not logits.is_cuda or logits.dtype == torch.float64: - # Triton kernel computes in float32; fall back for float64 and CPU - return selective_log_softmax_original( - logits, index.squeeze(-1) if squeeze else index - ) - - V = logits.shape[-1] - K = index.shape[-1] - original_index_shape = index.shape - - try: - flat_logits = logits.view(-1, V) - except RuntimeError: - flat_logits = logits.reshape(-1, V).contiguous() - flat_index = index.reshape(-1, K).contiguous() - - BLOCK_V = 4096 - MAX_GRID = 8192 - K_BLOCK = max(1, triton.next_power_of_2(K)) - - output = _SelectiveLogSoftmaxTriton.apply( - flat_logits, flat_index, K, K_BLOCK, V, BLOCK_V, MAX_GRID - ) - - if K_BLOCK != K: - output = output[:, :K] - - per_token_logps = output.to(logits.dtype).reshape(original_index_shape) - - if squeeze: - per_token_logps = per_token_logps.squeeze(-1) - - return per_token_logps diff --git a/src/axolotl/utils/mistral/mistral_tokenizer.py b/src/axolotl/utils/mistral/mistral_tokenizer.py index d90d6e5bb..bee6ca9e7 100644 --- a/src/axolotl/utils/mistral/mistral_tokenizer.py +++ b/src/axolotl/utils/mistral/mistral_tokenizer.py @@ -125,15 +125,20 @@ class HFMistralTokenizer(MistralCommonBackend): except (AttributeError, NotImplementedError): system_prompt = None - return normalizer._instruct_request_class( + request_class = normalizer._instruct_request_class + continuation_kwargs = {} + if "continue_final_message" in request_class.model_fields: + continuation_kwargs["continue_final_message"] = getattr( + request, "continue_final_message", False + ) + + return request_class( messages=messages, system_prompt=system_prompt, available_tools=request.tools, truncate_at_max_tokens=None, - continue_final_message=getattr( - request, "continue_final_message", False - ), settings=settings, + **continuation_kwargs, ) normalizer.from_chat_completion_request = from_chat_completion_request diff --git a/src/axolotl/utils/schemas/validation.py b/src/axolotl/utils/schemas/validation.py index 6abe53fe6..267a276c2 100644 --- a/src/axolotl/utils/schemas/validation.py +++ b/src/axolotl/utils/schemas/validation.py @@ -1735,12 +1735,6 @@ class SystemValidationMixin: checker = is_flash_attn_3_available else: checker = is_flash_attn_2_available - # TODO: remove when transformers 5.18.0 is released - from transformers.integrations.hub_kernels import ( - _FLASH_ATTN_KERNEL_VERSION_MAPPING, # pylint: disable=protected-access - ) - - _FLASH_ATTN_KERNEL_VERSION_MAPPING[2] = 3 with allow_all_hub_kernels(): available = checker(kernels_fallback_ok=True) reason = None diff --git a/tests/e2e/multigpu/solo/_ep_cpu_offload_grad_clip_worker.py b/tests/e2e/multigpu/solo/_ep_cpu_offload_grad_clip_worker.py index 389fdd848..93feece38 100644 --- a/tests/e2e/multigpu/solo/_ep_cpu_offload_grad_clip_worker.py +++ b/tests/e2e/multigpu/solo/_ep_cpu_offload_grad_clip_worker.py @@ -5,7 +5,6 @@ import torch.distributed as dist from torch.distributed.device_mesh import DeviceMesh from torch.distributed.tensor import DTensor, Replicate, Shard -from axolotl.monkeypatch.accelerate.parallelism_config import _ep_aware_clip_grad_norm from axolotl.utils.gradient_clipping import ( clip_grad_norm_ep_local_shards_, ep_local_parameter_ids, @@ -87,7 +86,7 @@ def main(): ) replica.grad = _dtensor(torch.tensor([7.0]), all_mesh, (Replicate(),), (1,)) plain_replica.grad = torch.tensor([6.0]) - gpu_path_norm = _ep_aware_clip_grad_norm( + gpu_path_norm = clip_grad_norm_ep_local_shards_( [expert, nonexpert, replica, plain_replica], torch.tensor(279.0).sqrt().item() / 3, ep_local_parameters={id(expert)}, @@ -121,7 +120,9 @@ def main(): torch.tensor([3.0 if dp == 0 else 4.0]), dp_mesh, (Shard(0),), (2,) ) torch.testing.assert_close( - _ep_aware_clip_grad_norm([subgroup], 10.0, global_mesh=mesh), + clip_grad_norm_ep_local_shards_( + [subgroup], 10.0, ep_local_parameters=set(), global_mesh=mesh + ), torch.tensor(5.0), ) dist.barrier() diff --git a/tests/integrations/conftest.py b/tests/integrations/conftest.py new file mode 100644 index 000000000..82ce613b3 --- /dev/null +++ b/tests/integrations/conftest.py @@ -0,0 +1,34 @@ +"""Fixtures shared by the expert-parallel integration tests.""" + +import pytest + + +@pytest.fixture +def fake_ep_sharder(monkeypatch): + """Run the real ``shard_expert_weights`` as one rank of a faked two-rank EP group. + + The collective scatter is replaced by a CPU slice; everything else the sharder + does to the module (counts, offsets, flags, DDP ignores) is the production code. + """ + from axolotl.integrations.expert_parallel import shard + + def run(model, rank, world_size=2): + monkeypatch.setattr(shard.dist, "get_world_size", lambda group=None: world_size) + monkeypatch.setattr(shard.dist, "get_rank", lambda group=None: rank) + monkeypatch.setattr( + shard.dist, + "all_gather_object", + lambda ranks, value: ranks.__setitem__( + slice(None), list(range(world_size)) + ), + ) + + def scatter_on_cpu(module, name, count, ranks): + shard._replace_with_slice( + module, name, ranks[rank] * count, (ranks[rank] + 1) * count + ) + + monkeypatch.setattr(shard, "_scatter_expert_from_rank0", scatter_on_cpu) + return shard.shard_expert_weights(model, object()) + + return run diff --git a/tests/integrations/test_expert_parallel.py b/tests/integrations/test_expert_parallel.py index 3569960f3..643a9e598 100644 --- a/tests/integrations/test_expert_parallel.py +++ b/tests/integrations/test_expert_parallel.py @@ -537,6 +537,25 @@ class TestCheckpointSaveScoping: # --------------------------------------------------------------------------- # +@pytest.mark.parametrize("rank", [0, 1]) +def test_sharding_marks_experts_parallel(fake_ep_sharder, rank): + model = torch.nn.Module() + model.experts = torch.nn.Module() + model.experts.gate_up_proj = torch.nn.Parameter(torch.randn(4, 8, 4)) + model.experts.down_proj = torch.nn.Parameter(torch.randn(4, 4, 4)) + model.experts.num_experts = 4 + model.experts._is_expert_parallel = False + model.dense = torch.nn.Linear(4, 4) + original = model.experts.gate_up_proj.detach().clone() + assert fake_ep_sharder(model, rank) == 1 + assert model.experts._is_expert_parallel is True + assert model.experts.num_experts == model.experts.num_local_experts == 2 + torch.testing.assert_close( + model.experts.gate_up_proj, original[rank * 2 : (rank + 1) * 2] + ) + assert not getattr(model.dense, "_is_expert_parallel", False) + + class TestShardingSingleRank: """At world_size=1, sharding is a no-op.""" @@ -558,11 +577,13 @@ class TestShardingSingleRank: n = shard_expert_weights(block, dist.group.WORLD) assert n == 0 assert tuple(block.experts.gate_up_proj.shape) == original_shape + assert not block.experts._is_expert_parallel def test_none_group_no_op(self): block = _build_qwen3moe_block(num_experts=16) n = shard_expert_weights(block, None) assert n == 0 + assert not block.experts._is_expert_parallel # --------------------------------------------------------------------------- # diff --git a/tests/integrations/test_expert_parallel_sentinel_kernels.py b/tests/integrations/test_expert_parallel_sentinel_kernels.py new file mode 100644 index 000000000..f730952f0 --- /dev/null +++ b/tests/integrations/test_expert_parallel_sentinel_kernels.py @@ -0,0 +1,174 @@ +"""Tests for transformers' local expert kernels under the EP sentinel routing contract.""" + +import copy + +import pytest +import torch +from transformers.integrations.moe import ( + batched_mm_experts_forward, + grouped_mm_experts_forward, +) + +from axolotl.integrations.expert_parallel import experts_fn +from axolotl.integrations.expert_parallel.experts_fn import ( + _eager_local, + _normalize_sentinels, +) + +E, K, H, INTER = 4, 2, 8, 16 +KINDS = ("qwen3_moe", "mixtral") +KERNELS = { + "grouped_mm": grouped_mm_experts_forward, + "batched_mm": batched_mm_experts_forward, +} + + +def _build_experts(kind): + torch.manual_seed(0) + if kind == "qwen3_moe": + from transformers.models.qwen3_moe.configuration_qwen3_moe import Qwen3MoeConfig + from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeExperts + + m = Qwen3MoeExperts( + Qwen3MoeConfig( + hidden_size=H, + moe_intermediate_size=INTER, + num_experts=E, + num_experts_per_tok=K, + norm_topk_prob=True, + ) + ) + else: + from transformers.models.mixtral.configuration_mixtral import MixtralConfig + from transformers.models.mixtral.modeling_mixtral import MixtralExperts + + m = MixtralExperts( + MixtralConfig( + hidden_size=H, + intermediate_size=INTER, + num_local_experts=E, + num_experts_per_tok=K, + ) + ) + with torch.no_grad(): + m.gate_up_proj.normal_(0, 0.2) + m.down_proj.normal_(0, 0.2) + return m + + +def _routing(num_local): + g = torch.Generator().manual_seed(7) + x = torch.randn(6, H, generator=g) + w_raw = torch.rand(6, K, generator=g) + if num_local == 4: + idx_raw = torch.tensor([[0, 1], [-1, 2], [-1, -1], [3, 0], [1, -1], [2, 3]]) + else: + idx_raw = torch.tensor([[0, 1], [-1, 1], [-1, -1], [1, 0], [0, -1], [1, 1]]) + return x, idx_raw, w_raw + + +def _prepare(kind, sharded, fake_ep_sharder=None): + m = _build_experts(kind) + if sharded: + holder = torch.nn.Module() + holder.experts = m + assert fake_ep_sharder(holder, rank=0) == 1 + assert m._is_expert_parallel is True + num_local = getattr(m, "num_local_experts", m.num_experts) + x, idx_raw, w_raw = _routing(num_local) + idx, w = _normalize_sentinels(idx_raw, w_raw, num_local) + return m, num_local, x, idx, w + + +@pytest.fixture(autouse=True) +def _restore_experts_fn_globals(): + prev = experts_fn.get_backend() + yield + experts_fn.set_backend(prev) + experts_fn.set_dispatch_chunks(1) + + +class TestSentinelKernelsWithFlag: + @pytest.mark.parametrize("sharded", (False, True)) + @pytest.mark.parametrize("kind", KINDS) + @pytest.mark.parametrize("kernel", KERNELS) + def test_matches_eager_and_grads_finite( + self, fake_ep_sharder, kernel, kind, sharded + ): + m, num_local, x, idx, w = _prepare(kind, sharded, fake_ep_sharder) + if not sharded: + m._is_expert_parallel = True + assert m._is_expert_parallel is True + sentinel = idx == num_local + assert int(idx.max()) == num_local + assert (w[sentinel] == 0).all() + + ref = copy.deepcopy(m) + xr = x.clone().requires_grad_(True) + wr = w.clone().requires_grad_(True) + yr = _eager_local(ref, xr, idx, wr) + + xk = x.clone().requires_grad_(True) + wk = w.clone().requires_grad_(True) + yk = KERNELS[kernel](m, xk, idx, wk) + + assert torch.isfinite(yk).all() + assert torch.allclose(yk, yr, atol=1e-6, rtol=1e-5) + assert torch.equal(yk[2], torch.zeros(H)) + + gout = torch.randn(6, H, generator=torch.Generator().manual_seed(21)) + (yk * gout).sum().backward() + (yr * gout).sum().backward() + pairs = [ + (xk.grad, xr.grad), + (wk.grad, wr.grad), + (m.gate_up_proj.grad, ref.gate_up_proj.grad), + (m.down_proj.grad, ref.down_proj.grad), + ] + for got, want in pairs: + assert got is not None and want is not None + assert torch.isfinite(got).all() + assert torch.allclose(got, want, atol=1e-6) + assert torch.equal(wk.grad[sentinel], torch.zeros_like(wk.grad[sentinel])) + + @pytest.mark.parametrize("kind", KINDS) + @pytest.mark.parametrize("kernel", KERNELS) + def test_caller_routing_not_mutated(self, kernel, kind): + m, _, x, idx, w = _prepare(kind, sharded=False) + m._is_expert_parallel = True + idx_copy = idx.clone() + w_copy = w.clone() + KERNELS[kernel](m, x, idx, w) + assert torch.equal(idx, idx_copy) + assert torch.equal(w, w_copy) + + +class TestSentinelKernelsWithoutFlag: + """grouped_mm without the flag leaves its sentinel-tail rows uninitialized upstream, so + its output is not asserted either way.""" + + @pytest.mark.parametrize("kind", KINDS) + def test_batched_mm_rejects_sentinel_without_flag(self, kind): + m, _, x, idx, w = _prepare(kind, sharded=False) + m._is_expert_parallel = False + with pytest.raises(IndexError): + batched_mm_experts_forward(m, x, idx, w) + + @pytest.mark.parametrize("kind", KINDS) + def test_ep_forward_sets_flag_before_local_kernel(self, kind): + m = _build_experts(kind) + assert m._is_expert_parallel is False + experts_fn.set_backend("torch") + x, idx_raw, w_raw = _routing(E) + y = experts_fn._ep_forward( + m, x, idx_raw, w_raw, local="batched_mm", backend="torch" + ) + assert m._is_expert_parallel is True + y_eager = experts_fn._ep_forward( + m, x, idx_raw, w_raw, local="eager", backend="torch" + ) + assert torch.allclose(y, y_eager, atol=1e-6) + y_grouped = experts_fn._ep_forward( + m, x, idx_raw, w_raw, local="grouped_mm", backend="torch" + ) + assert torch.allclose(y_grouped, y_eager, atol=1e-6) diff --git a/tests/integrations/test_expert_parallel_torch.py b/tests/integrations/test_expert_parallel_torch.py index c81da618b..ab28ccfc5 100644 --- a/tests/integrations/test_expert_parallel_torch.py +++ b/tests/integrations/test_expert_parallel_torch.py @@ -638,10 +638,12 @@ def _local_impl_checks(rank, world_size): ep = copy.deepcopy(full) s = _shard_experts(ep, rank, world_size) + assert ep._is_expert_parallel is False xe, we = x.clone().requires_grad_(True), w.clone().requires_grad_(True) experts_fn.set_dispatch_chunks(chunks) try: ye = experts_fn._ep_forward(ep, xe, idx, we, local=local, backend="torch") + assert ep._is_expert_parallel is True finally: experts_fn.set_dispatch_chunks(1) (ye * gout).sum().backward() @@ -1215,70 +1217,6 @@ class TestGatherEpExpertsIntoStateDict: assert torch.equal(r0["down"][2:], torch.full((2, 4, 2), 2.0)) -class TestEpClipGradNormPatchGate: - """The EP-aware clip engages for pure EP under FSDP2, not only for EP x dp_shard/cp.""" - - @staticmethod - def _fake_accelerator(ep, dp_shard, cp, fsdp2): - from types import SimpleNamespace - - return SimpleNamespace( - parallelism_config=( - None - if ep == "env" - else SimpleNamespace( - ep_enabled=ep, dp_shard_enabled=dp_shard, cp_enabled=cp - ) - ), - is_fsdp2=fsdp2, - unscale_gradients=lambda: None, - ) - - @pytest.mark.parametrize( - "ep,dp_shard,cp,fsdp2,expect_ep_aware", - [ - (True, False, False, True, True), - (True, True, False, False, True), - (True, False, True, False, True), - (True, False, False, False, False), - (False, False, False, True, False), - # pure EP: no ParallelismConfig at all, only the env var - ("env", False, False, True, True), - ("env", False, False, False, False), - ], - ) - def test_gate(self, monkeypatch, ep, dp_shard, cp, fsdp2, expect_ep_aware): - from accelerate import Accelerator - - from axolotl.monkeypatch.accelerate import parallelism_config as pc_mod - - if ep == "env": - monkeypatch.setenv("PARALLELISM_CONFIG_EP_SIZE", "2") - else: - monkeypatch.delenv("PARALLELISM_CONFIG_EP_SIZE", raising=False) - calls = [] - monkeypatch.setattr( - Accelerator, - "clip_grad_norm_", - lambda self, p, m, norm_type=2: calls.append("orig"), - ) - monkeypatch.setattr( - Accelerator, "_AXOLOTL_EP_CLIP_PATCHED", False, raising=False - ) - monkeypatch.setattr( - pc_mod, - "_ep_aware_clip_grad_norm", - lambda p, m, norm_type=2.0, **_kw: calls.append("ep_aware"), - ) - pc_mod.patch_clip_grad_norm_for_ep() - param = torch.nn.Parameter(torch.ones(3)) - param.grad = torch.ones(3) - Accelerator.clip_grad_norm_( - self._fake_accelerator(ep, dp_shard, cp, fsdp2), [param], 1.0 - ) - assert calls == (["ep_aware"] if expect_ep_aware else ["orig"]) - - def _composed_fsdp_worker(rank, world_size, port, q): """EP x dp_shard on a (2, 2) CPU mesh, EP x dp_replicate x dp_shard on (2, 2, 2), or dp_replicate x ep with no shard axis: expert grads must be the mean over every rank's diff --git a/tests/integrations/test_expert_parallel_transformers_contract.py b/tests/integrations/test_expert_parallel_transformers_contract.py new file mode 100644 index 000000000..3663e49f6 --- /dev/null +++ b/tests/integrations/test_expert_parallel_transformers_contract.py @@ -0,0 +1,327 @@ +"""Contract tests between the Expert-Parallel integration and transformers' MoE seam.""" + +import copy +import inspect +import os +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist +from transformers.distributed.tensor_parallel import ( + ALL_PARALLEL_STYLES, + EpRouterParallel, + MoEParamShard, +) +from transformers.integrations.moe import ( + ALL_EXPERTS_FUNCTIONS, + batched_mm_experts_forward, + grouped_mm_experts_forward, + use_experts_implementation, +) +from transformers.modeling_utils import PreTrainedModel + +from axolotl.integrations.expert_parallel import experts_fn +from axolotl.integrations.expert_parallel.experts_fn import ( + EXPERT_PARALLEL, + REGISTRY, + _decorator_experts_interface, + _normalize_sentinels, + register_all, +) +from axolotl.integrations.expert_parallel.shard import ( + _detect_experts_modules, + _is_param_wrapper, + _real_experts_base, +) + +E, K, H, INTER = 4, 2, 8, 16 +KINDS = ("qwen3_moe", "mixtral") + + +def _qwen3moe_config(): + from transformers.models.qwen3_moe.configuration_qwen3_moe import Qwen3MoeConfig + + return Qwen3MoeConfig( + hidden_size=H, + moe_intermediate_size=INTER, + num_experts=E, + num_experts_per_tok=K, + norm_topk_prob=True, + ) + + +def _mixtral_config(): + from transformers.models.mixtral.configuration_mixtral import MixtralConfig + + return MixtralConfig( + hidden_size=H, + intermediate_size=INTER, + num_local_experts=E, + num_experts_per_tok=K, + ) + + +def _build_experts(kind): + torch.manual_seed(0) + if kind == "qwen3_moe": + from transformers.models.qwen3_moe.modeling_qwen3_moe import ( + Qwen3MoeExperts as cls, + ) + + cfg = _qwen3moe_config() + else: + from transformers.models.mixtral.modeling_mixtral import MixtralExperts as cls + + cfg = _mixtral_config() + m = cls(cfg) + with torch.no_grad(): + m.gate_up_proj.normal_(0, 0.2) + m.down_proj.normal_(0, 0.2) + return m + + +def _build_qwen3moe_block(): + from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeSparseMoeBlock + + torch.manual_seed(0) + block = Qwen3MoeSparseMoeBlock(_qwen3moe_config()) + with torch.no_grad(): + block.experts.gate_up_proj.normal_(0, 0.2) + block.experts.down_proj.normal_(0, 0.2) + block.gate.weight.normal_(0, 1.0) + return block + + +def _unwrap(y): + return y[0] if isinstance(y, tuple) else y + + +class TestDecoratedExpertsAttributes: + @pytest.mark.parametrize("kind", KINDS) + def test_module_attributes(self, kind): + m = _build_experts(kind) + assert m.has_gate is True + assert m.has_bias is False + assert m.is_transposed is False + assert m.is_concatenated is True + assert m._is_expert_parallel is False + assert m.num_experts == E + assert callable(m.act_fn) + assert isinstance(m.gate_up_proj, torch.nn.Parameter) + assert tuple(m.gate_up_proj.shape) == (E, 2 * INTER, H) + assert isinstance(m.down_proj, torch.nn.Parameter) + assert tuple(m.down_proj.shape) == (E, H, INTER) + assert hasattr(m, "_apply_gate") + assert hasattr(m, "config") + assert getattr(m.config, "_experts_implementation", None) is None + assert not hasattr(m, "gate_up_proj_bias") + + def test_decorator_signature(self): + p = inspect.signature(use_experts_implementation).parameters + assert p["experts_interface"].default is ALL_EXPERTS_FUNCTIONS + assert p["is_concatenated"].default is True + assert p["is_transposed"].default is False + assert p["has_bias"].default is False + assert p["has_gate"].default is True + assert _decorator_experts_interface() is ALL_EXPERTS_FUNCTIONS + + def test_kernel_call_convention(self): + for fn in (grouped_mm_experts_forward, batched_mm_experts_forward): + assert list(inspect.signature(fn).parameters) == [ + "self", + "hidden_states", + "top_k_index", + "top_k_weights", + ] + assert ALL_EXPERTS_FUNCTIONS["grouped_mm"] is grouped_mm_experts_forward + assert ALL_EXPERTS_FUNCTIONS["batched_mm"] is batched_mm_experts_forward + + +class TestShardExpertWeightsOnDecoratedModule: + @pytest.mark.parametrize("rank", [0, 1]) + @pytest.mark.parametrize("kind", KINDS) + def test_sharded_module_matches_kernel_contract(self, fake_ep_sharder, kind, rank): + model = torch.nn.Module() + model.experts = _build_experts(kind) + model.dense = torch.nn.Linear(4, 4) + orig_gu = model.experts.gate_up_proj.detach().clone() + orig_dn = model.experts.down_proj.detach().clone() + assert fake_ep_sharder(model, rank) == 1 + + m = model.experts + e_local = E // 2 + assert m.num_experts == m.num_local_experts == e_local + assert m.num_experts_global == E + assert m.local_expert_offset == rank * e_local + assert m._is_expert_parallel is True + assert model._ep_num_experts_global == E + lo, hi = rank * e_local, (rank + 1) * e_local + torch.testing.assert_close(m.gate_up_proj, orig_gu[lo:hi]) + torch.testing.assert_close(m.down_proj, orig_dn[lo:hi]) + assert isinstance(m.gate_up_proj, torch.nn.Parameter) + assert isinstance(m.down_proj, torch.nn.Parameter) + assert m.gate_up_proj.requires_grad and m.down_proj.requires_grad + assert m.has_gate is True + assert m.has_bias is False + assert m.is_transposed is False + assert m.is_concatenated is True + assert model._ddp_params_and_buffers_to_ignore == [ + "experts.gate_up_proj", + "experts.down_proj", + ] + assert not getattr(model.dense, "_is_expert_parallel", False) + + idx, w = _normalize_sentinels( + torch.tensor([[0, 1], [-1, 1], [-1, -1]]), + torch.rand(3, 2), + m.num_local_experts, + ) + x = torch.randn(3, H) + y = grouped_mm_experts_forward(m, x, idx, w) + assert torch.isfinite(y).all() + assert torch.allclose(y, experts_fn._eager_local(m, x, idx, w), atol=1e-6) + + +class TestRegistrationAndDetection: + def test_registry_resolves_through_decorator_interface(self): + register_all() + assert set(REGISTRY) == {EXPERT_PARALLEL} + stub = PreTrainedModel.__new__(PreTrainedModel) + for name, fn in REGISTRY.items(): + assert ALL_EXPERTS_FUNCTIONS.get_interface(name, None) is fn + assert _decorator_experts_interface().get_interface(name, None) is fn + stub.config = SimpleNamespace(_experts_implementation=name) + assert ( + PreTrainedModel.get_correct_experts_implementation(stub, name) == name + ) + assert ( + experts_fn.resolve_local_implementation("eager") is experts_fn._eager_local + ) + assert ( + experts_fn.resolve_local_implementation("grouped_mm") + is grouped_mm_experts_forward + ) + assert ( + experts_fn.resolve_local_implementation("batched_mm") + is batched_mm_experts_forward + ) + assert "sonicmoe" in ALL_EXPERTS_FUNCTIONS + + @pytest.mark.parametrize("kind", KINDS) + def test_decorated_forward_dispatches_to_expert_parallel(self, kind): + register_all() + if kind == "qwen3_moe": + block = _build_qwen3moe_block() + inputs = (torch.randn(1, 5, H),) + else: + block = _build_experts("mixtral") + idx = torch.tensor([[0, 1], [2, 3], [1, 2], [3, 0], [0, 2]]) + inputs = (torch.randn(5, H), idx, torch.rand(5, K)) + + ref = copy.deepcopy(block) + ref_experts = ref.experts if hasattr(ref, "experts") else ref + ref_experts.config = copy.copy(ref_experts.config) + ref_experts.config._experts_implementation = "eager" + ep = copy.deepcopy(block) + ep_experts = ep.experts if hasattr(ep, "experts") else ep + ep_experts.config = copy.copy(ep_experts.config) + ep_experts.config._experts_implementation = EXPERT_PARALLEL + + prev = experts_fn.get_backend() + experts_fn.set_backend("torch") + try: + y_ref = _unwrap(ref(*inputs)) + y_ep = _unwrap(ep(*inputs)) + finally: + experts_fn.set_backend(prev) + + assert torch.allclose(y_ep, y_ref, atol=1e-6) + assert ep_experts._is_expert_parallel is True + assert ref_experts._is_expert_parallel is False + + def test_detects_decorated_experts_in_real_block(self): + block = _build_qwen3moe_block() + assert list(_detect_experts_modules(block)) == [("experts", block.experts)] + holder = torch.nn.Module() + holder.experts = _build_experts("mixtral") + assert list(_detect_experts_modules(holder)) == [("experts", holder.experts)] + + @pytest.mark.parametrize("kind", KINDS) + def test_detection_skips_real_peft_param_wrapper(self, kind): + from peft import LoraConfig + from peft.tuners.lora.layer import ParamWrapper + + holder = torch.nn.Module() + base = _build_experts(kind) + holder.experts = ParamWrapper( + base, + "default", + parameter_name="gate_up_proj", + config=LoraConfig(r=2, lora_alpha=4, lora_dropout=0.0), + r=2, + lora_alpha=4, + ) + assert _is_param_wrapper(holder.experts) + assert _real_experts_base(holder.experts) is base + assert list(_detect_experts_modules(holder)) == [("experts.base_layer", base)] + assert base.gate_up_proj.dim() == 3 + + +class TestTransformersExpertParallelPlan: + def test_plan_styles_exist(self): + assert "grouped_gemm" in ALL_PARALLEL_STYLES + assert "ep_router" in ALL_PARALLEL_STYLES + assert isinstance(ALL_PARALLEL_STYLES["grouped_gemm"], MoEParamShard) + assert ALL_PARALLEL_STYLES["grouped_gemm"].shards_expert_dim is True + assert isinstance(ALL_PARALLEL_STYLES["ep_router"], EpRouterParallel) + + def test_ep_router_sentinel_matches_normalize_sentinels(self): + mesh = SimpleNamespace( + get_local_rank=lambda: 1, + size=lambda: 2, + ndim=1, + get_group=lambda *a: None, + ) + router = SimpleNamespace(num_experts=E) + logits = torch.zeros(3, E) + scores = torch.tensor([[0.5, 0.5], [0.3, 0.7], [0.2, 0.8]]) + ind = torch.tensor([[0, 3], [1, 2], [2, 3]]) + _, s_tf, i_tf = EpRouterParallel().transform_output_post_forward( + router, (logits, scores, ind), mesh + ) + i_ax, s_ax = _normalize_sentinels( + torch.tensor([[-1, 1], [-1, 0], [0, 1]]), scores, E // 2 + ) + assert torch.equal(i_tf, i_ax) + assert torch.equal(i_tf, torch.tensor([[2, 1], [2, 0], [0, 1]])) + assert torch.equal(s_tf, s_ax) + assert int(i_tf.max()) == E // 2 + + +class TestMoEParamShardMarksExpertParallel: + def setup_method(self): + if not dist.is_initialized(): + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "29571") + os.environ.setdefault("RANK", "0") + os.environ.setdefault("WORLD_SIZE", "1") + dist.init_process_group(backend="gloo", rank=0, world_size=1) + + def teardown_method(self): + if dist.is_initialized(): + dist.destroy_process_group() + + @pytest.mark.parametrize("kind", KINDS) + def test_shard_param_sets_flag_and_local_count(self, kind): + from torch.distributed.device_mesh import init_device_mesh + from torch.distributed.tensor import DTensor + + mesh = init_device_mesh("cpu", (1,), mesh_dim_names=("ep",)) + m = _build_experts(kind) + ALL_PARALLEL_STYLES["grouped_gemm"].shard_param(m, "gate_up_proj", mesh) + assert m._is_expert_parallel is True + assert m.num_experts == E + assert isinstance(m.gate_up_proj, torch.nn.Parameter) + assert isinstance(m.gate_up_proj.data, DTensor) + assert tuple(m.gate_up_proj.shape) == (E, 2 * INTER, H) diff --git a/tests/kernels/test_op_registry.py b/tests/kernels/test_op_registry.py index 2b3da6039..39d962f0a 100644 --- a/tests/kernels/test_op_registry.py +++ b/tests/kernels/test_op_registry.py @@ -40,13 +40,6 @@ class TestOpsRegistered: assert torch.ops.axolotl.rms_norm_gated_fwd is not None assert torch.ops.axolotl.rms_norm_gated_bwd is not None - def test_trainer_utils_ops_exist(self): - import axolotl.monkeypatch.trainer.utils # noqa: F401 - - assert torch.ops.axolotl.selective_log_softmax_fwd is not None - assert torch.ops.axolotl.selective_log_softmax_bwd is not None - assert torch.ops.axolotl.entropy_from_logits is not None - def test_ebft_ops_exist(self): import axolotl.core.trainers.ebft.kernels # noqa: F401 @@ -123,25 +116,6 @@ class TestNewOpsDispatchVisibility: mod(x, gate=g) assert any("axolotl.rms_norm_gated_fwd" in s for s in rec.seen) - @requires_cuda - def test_selective_log_softmax_visible_to_dispatch_mode(self): - from axolotl.monkeypatch.trainer.utils import selective_log_softmax - - logits = torch.randn(2, 16, 512, device="cuda", dtype=torch.bfloat16) - index = torch.randint(0, 512, (2, 16), device="cuda") - with _OpRecorder() as rec: - selective_log_softmax(logits, index) - assert any("axolotl.selective_log_softmax_fwd" in s for s in rec.seen) - - @requires_cuda - def test_entropy_visible_to_dispatch_mode(self): - from axolotl.monkeypatch.trainer.utils import entropy_from_logits - - logits = torch.randn(2, 16, 512, device="cuda", dtype=torch.bfloat16) - with _OpRecorder() as rec: - entropy_from_logits(logits) - assert any("axolotl.entropy_from_logits" in s for s in rec.seen) - @requires_cuda def test_ebft_ops_visible_to_dispatch_mode(self): from axolotl.core.trainers.ebft.kernels import ( @@ -187,32 +161,6 @@ class TestNewOpsFakeTensor: assert y.shape == x.shape and y.dtype == x.dtype assert rstd.shape == (16,) and rstd.dtype == torch.float32 - @requires_cuda - def test_selective_log_softmax_fwd_fake_shapes(self): - from torch._subclasses.fake_tensor import FakeTensorMode - - import axolotl.monkeypatch.trainer.utils # noqa: F401 - - with FakeTensorMode(): - logits = torch.empty(32, 512, device="cuda", dtype=torch.bfloat16) - index = torch.empty(32, 1, device="cuda", dtype=torch.int64) - out, lse = torch.ops.axolotl.selective_log_softmax_fwd( - logits, index, 1, 1, 512, 4096, 8192 - ) - assert out.shape == (32, 1) and out.dtype == torch.float32 - assert lse.shape == (32,) and lse.dtype == torch.float32 - - @requires_cuda - def test_entropy_fake_shapes(self): - from torch._subclasses.fake_tensor import FakeTensorMode - - import axolotl.monkeypatch.trainer.utils # noqa: F401 - - with FakeTensorMode(): - logits = torch.empty(2, 16, 512, device="cuda", dtype=torch.bfloat16) - out = torch.ops.axolotl.entropy_from_logits(logits, 128) - assert out.shape == (32,) and out.dtype == torch.float32 - @requires_cuda def test_ebft_fake_shapes(self): from torch._subclasses.fake_tensor import FakeTensorMode diff --git a/tests/monkeypatch/test_torchao_tp_lora.py b/tests/monkeypatch/test_torchao_tp_lora.py index ad5269d32..9ec157324 100644 --- a/tests/monkeypatch/test_torchao_tp_lora.py +++ b/tests/monkeypatch/test_torchao_tp_lora.py @@ -111,6 +111,65 @@ def test_prepare_is_idempotent_and_uses_factor_gradient_hook(monkeypatch): assert len(model.lora_A["second"].weight._backward_hooks) == 1 +class DTensor: + """What PEFT leaves on the sharded factor under transformers' DTensor TP.""" + + def __init__(self, local, placement): + self._local = local + self.placements = (placement,) + self.requires_grad = True + + def to_local(self): + return self._local + + +def _peft_wrapped(model, factor_name, dim): + factor = model.__getattr__(factor_name)["default"] + local = factor.weight.detach().clone() + factor._parameters["weight"] = DTensor(local, Shard(dim)) + factor.forward = lambda *args, **kwargs: None + return local + + +@pytest.mark.parametrize( + "layout,factor_name,dim", [("colwise", "lora_B", 0), ("rowwise", "lora_A", 1)] +) +def test_prepare_localizes_peft_dtensor_factor(monkeypatch, layout, factor_name, dim): + model = _LoRA() + if layout == "rowwise": + object.__setattr__( + model.base_layer, + "weight", + SimpleNamespace( + placements=(Shard(1),), + shape=(4, 8), + device_mesh=SimpleNamespace(get_group=lambda: "tp"), + to_local=lambda: SimpleNamespace(shape=(4, 4)), + ), + ) + local = _peft_wrapped(model, factor_name, dim) + monkeypatch.setattr(torchao_tp_lora, "_is_native_nvfp4_dtensor", lambda _: True) + monkeypatch.setattr(torchao_tp_lora, "_lora_tp_plan", lambda *_: layout) + monkeypatch.setattr(torchao_tp_lora.dist, "get_world_size", lambda group: 1) + + assert torchao_tp_lora.prepare_native_nvfp4_tp_lora(model) + factor = model.__getattr__(factor_name)["default"] + assert isinstance(factor.weight, torch.nn.Parameter) + assert factor.weight.data_ptr() == local.data_ptr() + assert "forward" not in factor.__dict__ + assert factor.weight.requires_grad + + +def test_prepare_rejects_peft_factor_sharded_on_the_wrong_dim(monkeypatch): + model = _LoRA() + _peft_wrapped(model, "lora_A", 1) + monkeypatch.setattr(torchao_tp_lora, "_is_native_nvfp4_dtensor", lambda _: True) + monkeypatch.setattr(torchao_tp_lora, "_lora_tp_plan", lambda *_: "colwise") + + with pytest.raises(ValueError, match="expects lora_B sharded on dim 0"): + torchao_tp_lora.prepare_native_nvfp4_tp_lora(model) + + def test_prepare_rejects_unsynchronized_dropout(monkeypatch): model = _LoRA() model.lora_dropout["default"] = torch.nn.Dropout(0.1) diff --git a/tests/prompt_strategies/test_chat_templates_mistral.py b/tests/prompt_strategies/test_chat_templates_mistral.py index 956ce119a..bf43a608b 100644 --- a/tests/prompt_strategies/test_chat_templates_mistral.py +++ b/tests/prompt_strategies/test_chat_templates_mistral.py @@ -310,6 +310,7 @@ def test_mistral_chat_template( def test_mistral_normalizer_fills_instruct_request_defaults( magistral_tokenizer: "HFMistralTokenizer", + monkeypatch, ): from pydantic import BaseModel @@ -322,8 +323,10 @@ def test_mistral_normalizer_fills_instruct_request_defaults( def broken_from_chat_completion_request(_request): _MissingDefaults() - normalizer.from_chat_completion_request = broken_from_chat_completion_request - normalizer._axolotl_instruct_defaults_patched = False + monkeypatch.setattr( + normalizer, "from_chat_completion_request", broken_from_chat_completion_request + ) + monkeypatch.setattr(normalizer, "_axolotl_instruct_defaults_patched", False) magistral_tokenizer._patch_instruct_request_normalizer() input_ids = magistral_tokenizer.apply_chat_template( diff --git a/tests/test_distributed_parallel_grad_norm.py b/tests/test_distributed_parallel_grad_norm.py index d331ad0ce..2e204da3b 100644 --- a/tests/test_distributed_parallel_grad_norm.py +++ b/tests/test_distributed_parallel_grad_norm.py @@ -2,6 +2,8 @@ from types import SimpleNamespace +import pytest + def test_cpu_offload_telemetry_unscales_before_nonmutating_norm(monkeypatch): from torch.distributed.fsdp import CPUOffloadPolicy @@ -55,3 +57,61 @@ def test_existing_grad_norm_bypasses_cpu_offload_telemetry(): == 3.0 ) assert events == [] + + +def _ep_trainer(monkeypatch, parallelism, max_grad_norm=1.0): + from axolotl.core.trainers.mixins.distributed_parallel import ( + DistributedParallelMixin, + ) + from axolotl.utils import gradient_clipping + + events = [] + trainer = object.__new__(DistributedParallelMixin) + trainer.accelerator = SimpleNamespace( + state=SimpleNamespace(fsdp_plugin=SimpleNamespace(cpu_offload=None)), + parallelism_config=parallelism, + torch_device_mesh="mesh", + unscale_gradients=lambda: events.append("unscale"), + ) + trainer.args = SimpleNamespace(max_grad_norm=max_grad_norm) + monkeypatch.setattr(gradient_clipping, "ep_local_parameter_ids", lambda _: {1}) + + def clip(parameters, max_norm, *, ep_local_parameters, global_mesh): + events.append(("clip", max_norm, ep_local_parameters, global_mesh)) + return 5.0 + + def norm(parameters, *, ep_local_parameters, global_mesh): + events.append(("norm", ep_local_parameters, global_mesh)) + return 6.0 + + monkeypatch.setattr(gradient_clipping, "clip_grad_norm_ep_local_shards_", clip) + monkeypatch.setattr(gradient_clipping, "get_grad_norm_ep_local_shards_", norm) + return trainer, events + + +@pytest.mark.parametrize("via", ["parallelism_config", "env"]) +def test_expert_parallel_owns_clipping_without_cpu_offload(monkeypatch, via): + if via == "env": + monkeypatch.setenv("PARALLELISM_CONFIG_EP_SIZE", "2") + parallelism = None + else: + monkeypatch.delenv("PARALLELISM_CONFIG_EP_SIZE", raising=False) + parallelism = SimpleNamespace(ep_enabled=True) + trainer, events = _ep_trainer(monkeypatch, parallelism) + model = SimpleNamespace(parameters=lambda: [object()]) + + assert trainer._clip_grad_norm(model) == 5.0 + assert events == ["unscale", ("clip", 1.0, {1}, "mesh")] + events.clear() + assert trainer._get_grad_norm(model) == 6.0 + assert events == ["unscale", ("norm", {1}, "mesh")] + + +def test_dense_runs_keep_the_trainer_clipping(monkeypatch): + from transformers import Trainer + + monkeypatch.delenv("PARALLELISM_CONFIG_EP_SIZE", raising=False) + trainer, events = _ep_trainer(monkeypatch, SimpleNamespace(ep_enabled=False)) + monkeypatch.setattr(Trainer, "_clip_grad_norm", lambda self, model: "trainer") + assert trainer._clip_grad_norm(SimpleNamespace(parameters=lambda: [])) == "trainer" + assert events == [] diff --git a/tests/test_prm.py b/tests/test_prm.py new file mode 100644 index 000000000..a34bf7682 --- /dev/null +++ b/tests/test_prm.py @@ -0,0 +1,112 @@ +"""CPU regression coverage for the vendored PRM trainer.""" + +import numpy as np +import pytest +from datasets import Dataset +from tokenizers import Tokenizer +from tokenizers.models import WordLevel +from tokenizers.pre_tokenizers import Whitespace +from transformers import BertConfig, BertForTokenClassification, PreTrainedTokenizerFast + +from axolotl.core.trainers.prm.prm_trainer import PRMTrainer, compute_accuracy +from axolotl.core.trainers.trl import AxolotlPRMTrainer +from axolotl.core.training_args import AxolotlPRMConfig + + +@pytest.fixture +def prm_tokenizer(): + tokenizer = Tokenizer( + WordLevel( + {"[PAD]": 0, "[UNK]": 1, "question": 2, "good": 3, "bad": 4, "|": 5}, + unk_token="[UNK]", + ) + ) + tokenizer.pre_tokenizer = Whitespace() + return PreTrainedTokenizerFast( + tokenizer_object=tokenizer, pad_token="[PAD]", unk_token="[UNK]" + ) + + +@pytest.mark.parametrize( + "last_only,is_eval,labels", + [ + (False, False, [-100, -100, 1, -100, 0]), + (True, False, [-100, -100, -100, -100, 0]), + (True, True, [-100, -100, 1, -100, 0]), + ], +) +def test_step_labels(prm_tokenizer, last_only, is_eval, labels): + row = PRMTrainer.tokenize_row( + {"prompt": "question", "completions": ["good", "bad"], "labels": [True, False]}, + prm_tokenizer, + step_separator="|", + max_length=None, + max_completion_length=None, + train_on_last_step_only=last_only, + is_eval=is_eval, + ) + assert row == {"input_ids": [2, 3, 5, 4, 5], "labels": labels} + + +def test_accuracy_ignores_masked_tokens(): + logits = np.array([[[0, 1], [0, 1], [0, 1]]]) + labels = np.array([[-100, 1, 0]]) + assert compute_accuracy((logits, labels)) == {"accuracy": 0.5} + + +@pytest.mark.parametrize("pretokenized", [False, True]) +def test_axolotl_prm_train_and_evaluate(tmp_path, prm_tokenizer, pretokenized): + args = AxolotlPRMConfig( + output_dir=str(tmp_path), + use_cpu=True, + bf16=False, + gradient_checkpointing=False, + report_to="none", + max_steps=1, + per_device_train_batch_size=2, + per_device_eval_batch_size=2, + save_strategy="steps", + save_steps=1, + step_separator="|", + ) + model = BertForTokenClassification( + BertConfig( + vocab_size=len(prm_tokenizer), + hidden_size=16, + num_hidden_layers=1, + num_attention_heads=2, + intermediate_size=32, + num_labels=2, + ) + ) + if pretokenized: + dataset = Dataset.from_dict( + { + "input_ids": [[2, 3, 5], [2, 4, 5]], + "labels": [[-100, -100, 1], [-100, -100, 0]], + } + ) + else: + dataset = Dataset.from_dict( + { + "prompt": ["question"] * 2, + "completions": [["good"], ["bad"]], + "labels": [[True], [False]], + } + ) + trainer = AxolotlPRMTrainer( + model=model, + args=args, + processing_class=prm_tokenizer, + train_dataset=dataset, + eval_dataset=dataset, + ) + assert isinstance(trainer, PRMTrainer) + assert trainer.train_dataset[0]["labels"] == [-100, -100, 1] + assert np.isfinite(trainer.train().training_loss) + assert trainer.state.global_step == 1 + metrics = trainer.evaluate() + assert np.isfinite(metrics["eval_loss"]) + assert 0 <= metrics["eval_accuracy"] <= 1 + assert (tmp_path / "checkpoint-1" / "model.safetensors").is_file() + assert (tmp_path / "README.md").is_file() diff --git a/tests/test_triton_kernels.py b/tests/test_triton_kernels.py deleted file mode 100644 index ffc8de865..000000000 --- a/tests/test_triton_kernels.py +++ /dev/null @@ -1,481 +0,0 @@ -# Copyright 2026 Axolotl AI. 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. - -"""Unit tests for Triton kernels: entropy_from_logits and selective_log_softmax. - -Adapted from harness/test_entropy.py and harness/test_selective_logsoftmax.py -into proper pytest tests, plus new OOB index safety tests. -""" - -import math - -import pytest -import torch -import torch.nn.functional as F - -pytestmark = pytest.mark.skipif( - not torch.cuda.is_available(), reason="CUDA required for Triton kernels" -) - - -# --------------------------------------------------------------------------- -# Reference implementations -# --------------------------------------------------------------------------- - - -def _ref_entropy(logits): - """Reference entropy via log_softmax (numerically stable).""" - logp = F.log_softmax(logits.float(), dim=-1) - return -(logp.exp() * logp).sum(dim=-1) - - -def _ref_selective_log_softmax(logits, index): - """Reference selective log softmax via PyTorch gather.""" - squeeze = index.ndim == logits.ndim - 1 - if squeeze: - index = index.unsqueeze(-1) - log_probs = F.log_softmax(logits.float(), dim=-1) - result = torch.gather(log_probs, dim=-1, index=index) - if squeeze: - result = result.squeeze(-1) - return result - - -# --------------------------------------------------------------------------- -# entropy_from_logits -# --------------------------------------------------------------------------- - - -class TestEntropyFromLogits: - @pytest.mark.parametrize( - "B,L", - [ - (1, 128), - (1, 2048), - (4, 512), - (8, 256), - (1, 1), - ], - ) - def test_correctness_various_shapes(self, B, L): - from axolotl.monkeypatch.trainer.utils import entropy_from_logits - - V = 1024 - torch.manual_seed(42) - logits = torch.randn(B, L, V, device="cuda", dtype=torch.float32) - result = entropy_from_logits(logits) - expected = _ref_entropy(logits) - assert result.shape == (B, L) - torch.testing.assert_close(result, expected, atol=1e-4, rtol=1e-4) - - def test_2d_input(self): - from axolotl.monkeypatch.trainer.utils import entropy_from_logits - - logits = torch.randn(16, 256, device="cuda", dtype=torch.float32) - result = entropy_from_logits(logits) - expected = _ref_entropy(logits) - assert result.shape == (16,) - torch.testing.assert_close(result, expected, atol=1e-4, rtol=1e-4) - - def test_large_vocab(self): - from axolotl.monkeypatch.trainer.utils import entropy_from_logits - - V = 32000 - logits = torch.randn(2, V, device="cuda", dtype=torch.float32) - result = entropy_from_logits(logits) - expected = _ref_entropy(logits) - torch.testing.assert_close(result, expected, atol=1e-4, rtol=1e-4) - - def test_uniform_distribution(self): - """Uniform logits -> entropy = log(V).""" - from axolotl.monkeypatch.trainer.utils import entropy_from_logits - - V = 1024 - logits = torch.zeros(2, V, device="cuda", dtype=torch.float32) - result = entropy_from_logits(logits) - expected_val = math.log(V) - torch.testing.assert_close( - result, - torch.full((2,), expected_val, device="cuda", dtype=torch.float32), - atol=1e-4, - rtol=1e-4, - ) - - def test_peaked_distribution(self): - """One-hot-like logits -> entropy near 0.""" - from axolotl.monkeypatch.trainer.utils import entropy_from_logits - - logits = torch.full((2, 128), -100.0, device="cuda", dtype=torch.float32) - logits[:, 0] = 100.0 - result = entropy_from_logits(logits) - assert (result < 1e-3).all() - - def test_bfloat16(self): - from axolotl.monkeypatch.trainer.utils import entropy_from_logits - - logits = torch.randn(4, 256, device="cuda", dtype=torch.bfloat16) - result = entropy_from_logits(logits) - expected = _ref_entropy(logits.float()) - assert result.dtype == torch.bfloat16 - torch.testing.assert_close(result.float(), expected, atol=5e-2, rtol=5e-2) - - def test_float16(self): - from axolotl.monkeypatch.trainer.utils import entropy_from_logits - - logits = torch.randn(4, 256, device="cuda", dtype=torch.float16) - result = entropy_from_logits(logits) - expected = _ref_entropy(logits.float()) - assert result.dtype == torch.float16 - torch.testing.assert_close(result.float(), expected, atol=5e-2, rtol=5e-2) - - def test_non_contiguous_3d_transpose(self): - """Non-contiguous 3D tensor via transpose(0,1).""" - from axolotl.monkeypatch.trainer.utils import entropy_from_logits - - V = 256 - raw = torch.randn(32, 4, V, device="cuda", dtype=torch.float32) - logits = raw.transpose(0, 1) # (4, 32, V) non-contiguous - assert not logits.is_contiguous() - result = entropy_from_logits(logits) - expected = _ref_entropy(logits) - torch.testing.assert_close(result, expected, atol=1e-4, rtol=1e-4) - - def test_non_contiguous_3d_slice(self): - """Non-contiguous 3D tensor via batch slicing.""" - from axolotl.monkeypatch.trainer.utils import entropy_from_logits - - V = 256 - raw = torch.randn(8, 32, V, device="cuda", dtype=torch.float32) - logits = raw[::2] # (4, 32, V) non-contiguous - assert not logits.is_contiguous() - result = entropy_from_logits(logits) - expected = _ref_entropy(logits) - torch.testing.assert_close(result, expected, atol=1e-4, rtol=1e-4) - - def test_many_rows_beyond_max_grid(self): - """More rows than MAX_GRID (8192) to test chunked dispatch.""" - from axolotl.monkeypatch.trainer.utils import entropy_from_logits - - logits = torch.randn(10000, 128, device="cuda", dtype=torch.float32) - result = entropy_from_logits(logits) - expected = _ref_entropy(logits) - torch.testing.assert_close(result, expected, atol=1e-4, rtol=1e-4) - - def test_entropy_non_negative(self): - from axolotl.monkeypatch.trainer.utils import entropy_from_logits - - logits = torch.randn(32, 512, device="cuda", dtype=torch.float32) - result = entropy_from_logits(logits) - assert (result >= -1e-5).all(), f"Negative entropy: {result.min()}" - - -# --------------------------------------------------------------------------- -# selective_log_softmax — forward correctness -# --------------------------------------------------------------------------- - - -class TestSelectiveLogSoftmax: - @pytest.mark.parametrize( - "B,L,K", - [ - (1, 128, 1), - (4, 512, 1), - (8, 256, 1), - (4, 256, 4), - (4, 256, 7), - (15, 129, 1), # non-power-of-2 - ], - ) - def test_correctness_various_shapes(self, B, L, K): - from axolotl.monkeypatch.trainer.utils import selective_log_softmax - - V = 1024 - torch.manual_seed(42) - logits = torch.randn(B, L, V, device="cuda", dtype=torch.float32) - if K == 1: - index = torch.randint(0, V, (B, L), device="cuda") - else: - index = torch.randint(0, V, (B, L, K), device="cuda") - result = selective_log_softmax(logits, index) - expected = _ref_selective_log_softmax(logits, index) - torch.testing.assert_close(result, expected, atol=1e-4, rtol=1e-4) - - def test_squeezed_index(self): - """Index with ndim == logits.ndim - 1 triggers squeeze path.""" - from axolotl.monkeypatch.trainer.utils import selective_log_softmax - - V = 256 - logits = torch.randn(8, V, device="cuda", dtype=torch.float32) - index = torch.randint(0, V, (8,), device="cuda") - result = selective_log_softmax(logits, index) - expected = _ref_selective_log_softmax(logits, index) - assert result.shape == (8,) - torch.testing.assert_close(result, expected, atol=1e-4, rtol=1e-4) - - def test_large_vocab(self): - from axolotl.monkeypatch.trainer.utils import selective_log_softmax - - V = 32000 - logits = torch.randn(2, V, device="cuda", dtype=torch.float32) - index = torch.randint(0, V, (2, 1), device="cuda") - result = selective_log_softmax(logits, index) - expected = _ref_selective_log_softmax(logits, index) - torch.testing.assert_close(result, expected, atol=1e-4, rtol=1e-4) - - def test_bfloat16(self): - from axolotl.monkeypatch.trainer.utils import selective_log_softmax - - V = 1024 - torch.manual_seed(42) - logits = torch.randn(4, 128, V, device="cuda", dtype=torch.bfloat16) - index = torch.randint(0, V, (4, 128), device="cuda") - result = selective_log_softmax(logits, index) - expected = _ref_selective_log_softmax(logits.float(), index) - assert result.dtype == torch.bfloat16 - torch.testing.assert_close(result.float(), expected, atol=0.1, rtol=0.1) - - def test_fp32_tight_tolerance(self): - from axolotl.monkeypatch.trainer.utils import selective_log_softmax - - V = 1024 - torch.manual_seed(42) - logits = torch.randn(2, 256, V, device="cuda", dtype=torch.float32) - index = torch.randint(0, V, (2, 256), device="cuda") - result = selective_log_softmax(logits, index) - expected = _ref_selective_log_softmax(logits, index) - torch.testing.assert_close(result, expected, atol=1e-5, rtol=1e-5) - - def test_all_same_index(self): - from axolotl.monkeypatch.trainer.utils import selective_log_softmax - - V = 128 - logits = torch.randn(8, V, device="cuda", dtype=torch.float32) - index = torch.zeros(8, 1, device="cuda", dtype=torch.long) - result = selective_log_softmax(logits, index) - expected = _ref_selective_log_softmax(logits, index) - torch.testing.assert_close(result, expected, atol=1e-4, rtol=1e-4) - - def test_last_index(self): - from axolotl.monkeypatch.trainer.utils import selective_log_softmax - - V = 128 - logits = torch.randn(8, V, device="cuda", dtype=torch.float32) - index = torch.full((8, 1), V - 1, device="cuda", dtype=torch.long) - result = selective_log_softmax(logits, index) - expected = _ref_selective_log_softmax(logits, index) - torch.testing.assert_close(result, expected, atol=1e-4, rtol=1e-4) - - def test_output_always_nonpositive(self): - """Log softmax values should always be <= 0.""" - from axolotl.monkeypatch.trainer.utils import selective_log_softmax - - V = 256 - logits = torch.randn(32, V, device="cuda", dtype=torch.float32) - index = torch.randint(0, V, (32, 1), device="cuda") - result = selective_log_softmax(logits, index) - assert (result <= 1e-5).all(), f"Positive log-prob: {result.max()}" - - def test_many_rows_beyond_max_grid(self): - from axolotl.monkeypatch.trainer.utils import selective_log_softmax - - V = 128 - logits = torch.randn(10000, V, device="cuda", dtype=torch.float32) - index = torch.randint(0, V, (10000, 1), device="cuda") - result = selective_log_softmax(logits, index) - expected = _ref_selective_log_softmax(logits, index) - torch.testing.assert_close(result, expected, atol=1e-4, rtol=1e-4) - - -# --------------------------------------------------------------------------- -# selective_log_softmax — backward / gradient correctness -# --------------------------------------------------------------------------- - - -class TestSelectiveLogSoftmaxBackward: - @pytest.mark.parametrize( - "B,L,V,K", - [ - (2, 16, 64, 1), - (2, 16, 64, 4), - (1, 8, 128, 1), - (2, 8, 128, 7), - ], - ) - def test_gradient_matches_reference(self, B, L, V, K): - from axolotl.monkeypatch.trainer.utils import selective_log_softmax - - torch.manual_seed(42) - logits_ref = torch.randn( - B, L, V, device="cuda", dtype=torch.float32, requires_grad=True - ) - logits_tri = logits_ref.detach().clone().requires_grad_(True) - - if K == 1: - index = torch.randint(0, V, (B, L), device="cuda") - else: - index = torch.randint(0, V, (B, L, K), device="cuda") - - ref_out = _ref_selective_log_softmax(logits_ref, index) - tri_out = selective_log_softmax(logits_tri, index) - - ref_out.sum().backward() - tri_out.sum().backward() - - torch.testing.assert_close( - logits_tri.grad, logits_ref.grad, atol=1e-5, rtol=1e-5 - ) - - def test_gradient_bfloat16_full_vocab(self): - from axolotl.monkeypatch.trainer.utils import selective_log_softmax - - V = 4096 - torch.manual_seed(42) - logits_ref = torch.randn( - 2, 64, V, device="cuda", dtype=torch.bfloat16, requires_grad=True - ) - logits_tri = logits_ref.detach().clone().requires_grad_(True) - index = torch.randint(0, V, (2, 64), device="cuda") - - _ref_selective_log_softmax(logits_ref, index).sum().backward() - selective_log_softmax(logits_tri, index).sum().backward() - - torch.testing.assert_close( - logits_tri.grad.float(), logits_ref.grad.float(), atol=0.1, rtol=0.1 - ) - - def test_gradient_k1_squeezed(self): - """Gradient with squeezed (1D) index.""" - from axolotl.monkeypatch.trainer.utils import selective_log_softmax - - V = 256 - logits = torch.randn( - 8, V, device="cuda", dtype=torch.float32, requires_grad=True - ) - index = torch.randint(0, V, (8,), device="cuda") - - result = selective_log_softmax(logits, index) - result.sum().backward() - triton_grad = logits.grad.clone() - - logits.grad = None - ref = torch.gather( - F.log_softmax(logits, dim=-1), dim=-1, index=index.unsqueeze(-1) - ).squeeze(-1) - ref.sum().backward() - - torch.testing.assert_close(triton_grad, logits.grad, atol=1e-4, rtol=1e-4) - - -# --------------------------------------------------------------------------- -# selective_log_softmax — out-of-bounds index safety -# --------------------------------------------------------------------------- - - -class TestSelectiveLogSoftmaxOOBSafety: - """Verify that out-of-range indices don't crash or corrupt valid results.""" - - def test_negative_indices_no_crash(self): - from axolotl.monkeypatch.trainer.utils import selective_log_softmax - - V = 128 - logits = torch.randn(4, V, device="cuda", dtype=torch.float32) - index = torch.tensor( - [[-1], [0], [V - 1], [-5]], device="cuda", dtype=torch.long - ) - result = selective_log_softmax(logits, index) - assert result.shape == (4, 1) - # Valid rows should be finite and match reference - valid_idx = torch.tensor([[0], [V - 1]], device="cuda", dtype=torch.long) - valid_logits = logits[1:3] - expected = _ref_selective_log_softmax(valid_logits, valid_idx) - torch.testing.assert_close(result[1:3], expected, atol=1e-4, rtol=1e-4) - - def test_index_exceeds_vocab_no_crash(self): - from axolotl.monkeypatch.trainer.utils import selective_log_softmax - - V = 128 - logits = torch.randn(4, V, device="cuda", dtype=torch.float32) - index = torch.tensor( - [[0], [V], [V + 100], [V - 1]], device="cuda", dtype=torch.long - ) - result = selective_log_softmax(logits, index) - assert result.shape == (4, 1) - # Valid rows (0 and 3) should match reference - for row_idx, idx_val in [(0, 0), (3, V - 1)]: - ref = _ref_selective_log_softmax( - logits[row_idx : row_idx + 1], - torch.tensor([[idx_val]], device="cuda", dtype=torch.long), - ) - torch.testing.assert_close( - result[row_idx : row_idx + 1], ref, atol=1e-4, rtol=1e-4 - ) - - def test_mixed_valid_invalid_multi_index(self): - from axolotl.monkeypatch.trainer.utils import selective_log_softmax - - V = 256 - K = 3 - logits = torch.randn(4, V, device="cuda", dtype=torch.float32) - index = torch.tensor( - [ - [0, 10, -1], # last invalid - [V, 5, 100], # first invalid - [50, 60, 70], # all valid - [-1, V + 1, -100], # all invalid - ], - device="cuda", - dtype=torch.long, - ) - result = selective_log_softmax(logits, index) - assert result.shape == (4, K) - # Row 2 (all valid) must match reference exactly - valid_index = torch.tensor([[50, 60, 70]], device="cuda", dtype=torch.long) - expected = _ref_selective_log_softmax(logits[2:3], valid_index) - torch.testing.assert_close(result[2:3], expected, atol=1e-4, rtol=1e-4) - - def test_oob_backward_no_crash(self): - """Backward with OOB indices should not crash and grads should be finite.""" - from axolotl.monkeypatch.trainer.utils import selective_log_softmax - - V = 128 - logits = torch.randn( - 4, V, device="cuda", dtype=torch.float32, requires_grad=True - ) - index = torch.tensor( - [[-1], [0], [V + 10], [V - 1]], device="cuda", dtype=torch.long - ) - result = selective_log_softmax(logits, index) - result.sum().backward() - assert logits.grad is not None - assert torch.isfinite(logits.grad).all() - - def test_oob_backward_valid_rows_correct(self): - """Gradients for valid-index rows should match reference even when other rows have OOB.""" - from axolotl.monkeypatch.trainer.utils import selective_log_softmax - - V = 128 - logits = torch.randn( - 4, V, device="cuda", dtype=torch.float32, requires_grad=True - ) - # Row 0: invalid, Row 1: valid, Row 2: invalid, Row 3: valid - index = torch.tensor( - [[-1], [42], [V + 5], [100]], device="cuda", dtype=torch.long - ) - result = selective_log_softmax(logits, index) - result.sum().backward() - - # Compute reference gradient for valid rows only - logits_ref = logits.detach().clone().requires_grad_(True) - valid_rows = [1, 3] - valid_indices = [42, 100] - for r, idx in zip(valid_rows, valid_indices, strict=True): - ref_lp = F.log_softmax(logits_ref[r : r + 1], dim=-1) - ref_val = ref_lp[0, idx] - ref_val.backward(retain_graph=True) - - for r in valid_rows: - torch.testing.assert_close( - logits.grad[r], logits_ref.grad[r], atol=1e-4, rtol=1e-4 - ) diff --git a/tests/utils/test_mistral_tokenizer.py b/tests/utils/test_mistral_tokenizer.py index 869bfda17..a3e0393e9 100644 --- a/tests/utils/test_mistral_tokenizer.py +++ b/tests/utils/test_mistral_tokenizer.py @@ -1,12 +1,74 @@ """Tests for the mistral-common tokenizer wrapper.""" +from types import SimpleNamespace from unittest.mock import patch import pytest +from mistral_common.protocol.instruct.messages import UserMessage +from mistral_common.protocol.instruct.normalize import InstructRequestNormalizerV7 +from mistral_common.protocol.instruct.request import ChatCompletionRequest +from pydantic import BaseModel, ValidationError, create_model from axolotl.utils.mistral.mistral_tokenizer import HFMistralTokenizer +@pytest.mark.parametrize("legacy_continuation", [False, True]) +@pytest.mark.parametrize("continue_final_message", [False, True]) +def test_normalizer_defaults_match_request_schema( + legacy_continuation, continue_final_message +): + normalizer = InstructRequestNormalizerV7.normalizer() + if legacy_continuation: + normalizer._instruct_request_class = create_model( + "LegacyInstructRequest", + __base__=normalizer._instruct_request_class, + continue_final_message=(bool, ...), + ) + + class MissingDefaults(BaseModel): + truncate_at_max_tokens: int | None + continue_final_message: bool + + def broken_normalize(_request): + MissingDefaults.model_validate({}) + + normalizer.from_chat_completion_request = broken_normalize + tokenizer = object.__new__(HFMistralTokenizer) + tokenizer.tokenizer = SimpleNamespace(_instruct_request_normalizer=normalizer) + tokenizer._patch_instruct_request_normalizer() + request = SimpleNamespace( + messages=[UserMessage(content="Hello")], + tools=None, + continue_final_message=continue_final_message, + ) + result = normalizer.from_chat_completion_request(request) + assert result.messages == request.messages + assert result.truncate_at_max_tokens is None + if legacy_continuation: + assert result.continue_final_message is continue_final_message + else: + assert "continue_final_message" not in result.model_dump() + + +def test_normalizer_preserves_unrelated_validation_errors(): + normalizer = InstructRequestNormalizerV7.normalizer() + + class RequiredMessages(BaseModel): + messages: list + + def broken_normalize(_request): + RequiredMessages.model_validate({}) + + normalizer.from_chat_completion_request = broken_normalize + tokenizer = object.__new__(HFMistralTokenizer) + tokenizer.tokenizer = SimpleNamespace(_instruct_request_normalizer=normalizer) + tokenizer._patch_instruct_request_normalizer() + with pytest.raises(ValidationError, match="messages"): + normalizer.from_chat_completion_request( + ChatCompletionRequest(messages=[UserMessage(content="Hello")]) + ) + + @pytest.fixture(name="captured_init") def captured_init_fixture(): """Capture the kwargs `from_pretrained` builds without constructing a tokenizer."""