mirror of
https://github.com/axolotl-ai-cloud/axolotl.git
synced 2026-10-11 21:59:55 +00:00
bump HF deps 20261001 (#4097)
* bump deps * bump hub too * bump hf-xet also * vendor deprecated PRM trainer * fix mistral tokenizer api change * remove override patch for fa2 * fix EP handling * Add contract tests for the expert-parallel seam with transformers 5.18 Pin every attribute, kernel signature, registry hook, and plan style that the expert-parallel integration borrows from transformers, on real decorated Qwen3-MoE and Mixtral experts modules, so the next dependency bump fails in a named CPU test. Cover the sentinel routing contract for grouped_mm and batched_mm with _is_expert_parallel set, and the IndexError batched_mm raises without it. * Localize PEFT's DTensor LoRA factors in the native NVFP4 TP path PEFT 0.21.1 wraps the sharded factor as a DTensor and installs transformers' tensor-parallel forward on it when the base layer carries a TP plan. The native NVFP4 TP LoRA forward computes with local factors and does its own reductions, so unwrap the factors and drop that forward before validating them. * Drop the Triton selective_log_softmax override in favour of trl's kernel trl 1.14.1 ships a fused log-softmax and entropy kernel that matches or beats the axolotl Triton version at every shape the RL trainers use, and takes the temperature and row_mask arguments the trainers now pass. * Own EP gradient clipping in the trainer instead of patching accelerate transformers 5.18 routes clipping through a per-mesh path whenever the gradients do not share one device mesh, which pure EP never does: experts live on their own mesh and the dense weights on the EP mesh. The per-mesh norm never sums the experts held by other ranks, so the global norm is too small and the clipped update lands 5-7% too large. The trainer's _clip_grad_norm/_get_grad_norm hooks now compute the ownership-filtered EP norm directly when EP is enabled, so no accelerate clip_grad_norm_ patch is needed. * Drop the removed monkeypatch.trainer.utils module from the API reference
This commit is contained in:
Wing Lian
@@ -1,3 +1,3 @@
|
||||
[bandit]
|
||||
exclude = tests
|
||||
exclude = tests,src/axolotl/core/trainers/prm
|
||||
skips = B101,B615,B102,B110
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
+15
-11
@@ -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"]
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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.
|
||||
@@ -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.
|
||||
@@ -0,0 +1,5 @@
|
||||
"""PRM components vendored from Hugging Face TRL; see NOTICE and LICENSE."""
|
||||
|
||||
from .prm_config import PRMConfig
|
||||
|
||||
__all__ = ["PRMConfig"]
|
||||
@@ -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."},
|
||||
)
|
||||
@@ -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)
|
||||
@@ -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"]
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -1,3 +1 @@
|
||||
from .utils import entropy_from_logits, selective_log_softmax
|
||||
|
||||
__all__ = ["entropy_from_logits", "selective_log_softmax"]
|
||||
"""Trainer-side monkeypatches."""
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
@@ -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)
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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 == []
|
||||
|
||||
@@ -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()
|
||||
@@ -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
|
||||
)
|
||||
@@ -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."""
|
||||
|
||||
Reference in New Issue
Block a user