.py`, and declare it on that instance.
+ """
+
+ # Class instance object, so that a call to `register` can be reflected into all other files correctly, even if
+ # a new instance is created (in order to locally override a given function)
+ _global_mapping = {
+ "flash_attention_2": flash_attention_forward,
+ "flex_attention": flex_attention_forward,
+ "sdpa": sdpa_attention_forward,
+ }
+
+ def __init__(self):
+ self._local_mapping = {}
+
+ def __getitem__(self, key):
+ # First check if instance has a local override
+ if key in self._local_mapping:
+ return self._local_mapping[key]
+ return self._global_mapping[key]
+
+ def __setitem__(self, key, value):
+ # Allow local update of the default functions without impacting other instances
+ self._local_mapping.update({key: value})
+
+ def __delitem__(self, key):
+ del self._local_mapping[key]
+
+ def __iter__(self):
+ # Ensure we use all keys, with the overwritten ones on top
+ return iter({**self._global_mapping, **self._local_mapping})
+
+ def __len__(self):
+ return len(self._global_mapping.keys() | self._local_mapping.keys())
+
+ @classmethod
+ def register(cls, key: str, value: Callable):
+ cls._global_mapping.update({key: value})
+
+ def valid_keys(self) -> List[str]:
+ return list(self.keys())
+
+
+# Global AttentionInterface shared by all models which do not need to overwrite any of the existing ones
+ALL_ATTENTION_FUNCTIONS: AttentionInterface = AttentionInterface()
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/notebook.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/notebook.py
new file mode 100644
index 0000000000000000000000000000000000000000..22a44d858e731f9bd4067d373165b0b376f5191b
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/notebook.py
@@ -0,0 +1,383 @@
+# Copyright 2020 Hugging Face
+#
+# 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 os
+import re
+import time
+from typing import Optional
+
+import IPython.display as disp
+
+from ..trainer_callback import TrainerCallback
+from ..trainer_utils import IntervalStrategy, has_length
+
+
+def format_time(t):
+ "Format `t` (in seconds) to (h):mm:ss"
+ t = int(t)
+ h, m, s = t // 3600, (t // 60) % 60, t % 60
+ return f"{h}:{m:02d}:{s:02d}" if h != 0 else f"{m:02d}:{s:02d}"
+
+
+def html_progress_bar(value, total, prefix, label, width=300):
+ # docstyle-ignore
+ return f"""
+
+ """
+
+
+def text_to_html_table(items):
+ "Put the texts in `items` in an HTML table."
+ html_code = """\n"""
+ html_code += """ \n \n"""
+ for i in items[0]:
+ html_code += f" | {i} | \n"
+ html_code += "
\n \n \n"
+ for line in items[1:]:
+ html_code += " \n"
+ for elt in line:
+ elt = f"{elt:.6f}" if isinstance(elt, float) else str(elt)
+ html_code += f" | {elt} | \n"
+ html_code += "
\n"
+ html_code += " \n
"
+ return html_code
+
+
+class NotebookProgressBar:
+ """
+ A progress par for display in a notebook.
+
+ Class attributes (overridden by derived classes)
+
+ - **warmup** (`int`) -- The number of iterations to do at the beginning while ignoring `update_every`.
+ - **update_every** (`float`) -- Since calling the time takes some time, we only do it every presumed
+ `update_every` seconds. The progress bar uses the average time passed up until now to guess the next value
+ for which it will call the update.
+
+ Args:
+ total (`int`):
+ The total number of iterations to reach.
+ prefix (`str`, *optional*):
+ A prefix to add before the progress bar.
+ leave (`bool`, *optional*, defaults to `True`):
+ Whether or not to leave the progress bar once it's completed. You can always call the
+ [`~utils.notebook.NotebookProgressBar.close`] method to make the bar disappear.
+ parent ([`~notebook.NotebookTrainingTracker`], *optional*):
+ A parent object (like [`~utils.notebook.NotebookTrainingTracker`]) that spawns progress bars and handle
+ their display. If set, the object passed must have a `display()` method.
+ width (`int`, *optional*, defaults to 300):
+ The width (in pixels) that the bar will take.
+
+ Example:
+
+ ```python
+ import time
+
+ pbar = NotebookProgressBar(100)
+ for val in range(100):
+ pbar.update(val)
+ time.sleep(0.07)
+ pbar.update(100)
+ ```"""
+
+ warmup = 5
+ update_every = 0.2
+
+ def __init__(
+ self,
+ total: int,
+ prefix: Optional[str] = None,
+ leave: bool = True,
+ parent: Optional["NotebookTrainingTracker"] = None,
+ width: int = 300,
+ ):
+ self.total = total
+ self.prefix = "" if prefix is None else prefix
+ self.leave = leave
+ self.parent = parent
+ self.width = width
+ self.last_value = None
+ self.comment = None
+ self.output = None
+ self.value = None
+ self.label = None
+ if "VSCODE_PID" in os.environ:
+ self.update_every = 0.5 # Adjusted for smooth updated as html rending is slow on VS Code
+ # This is the only adjustment required to optimize training html rending
+
+ def update(self, value: int, force_update: bool = False, comment: Optional[str] = None):
+ """
+ The main method to update the progress bar to `value`.
+
+ Args:
+ value (`int`):
+ The value to use. Must be between 0 and `total`.
+ force_update (`bool`, *optional*, defaults to `False`):
+ Whether or not to force and update of the internal state and display (by default, the bar will wait for
+ `value` to reach the value it predicted corresponds to a time of more than the `update_every` attribute
+ since the last update to avoid adding boilerplate).
+ comment (`str`, *optional*):
+ A comment to add on the left of the progress bar.
+ """
+ self.value = value
+ if comment is not None:
+ self.comment = comment
+ if self.last_value is None:
+ self.start_time = self.last_time = time.time()
+ self.start_value = self.last_value = value
+ self.elapsed_time = self.predicted_remaining = None
+ self.first_calls = self.warmup
+ self.wait_for = 1
+ self.update_bar(value)
+ elif value <= self.last_value and not force_update:
+ return
+ elif force_update or self.first_calls > 0 or value >= min(self.last_value + self.wait_for, self.total):
+ if self.first_calls > 0:
+ self.first_calls -= 1
+ current_time = time.time()
+ self.elapsed_time = current_time - self.start_time
+ # We could have value = self.start_value if the update is called twixe with the same start value.
+ if value > self.start_value:
+ self.average_time_per_item = self.elapsed_time / (value - self.start_value)
+ else:
+ self.average_time_per_item = None
+ if value >= self.total:
+ value = self.total
+ self.predicted_remaining = None
+ if not self.leave:
+ self.close()
+ elif self.average_time_per_item is not None:
+ self.predicted_remaining = self.average_time_per_item * (self.total - value)
+ self.update_bar(value)
+ self.last_value = value
+ self.last_time = current_time
+ if (self.average_time_per_item is None) or (self.average_time_per_item == 0):
+ self.wait_for = 1
+ else:
+ self.wait_for = max(int(self.update_every / self.average_time_per_item), 1)
+
+ def update_bar(self, value, comment=None):
+ spaced_value = " " * (len(str(self.total)) - len(str(value))) + str(value)
+ if self.elapsed_time is None:
+ self.label = f"[{spaced_value}/{self.total} : < :"
+ elif self.predicted_remaining is None:
+ self.label = f"[{spaced_value}/{self.total} {format_time(self.elapsed_time)}"
+ else:
+ self.label = (
+ f"[{spaced_value}/{self.total} {format_time(self.elapsed_time)} <"
+ f" {format_time(self.predicted_remaining)}"
+ )
+ if self.average_time_per_item == 0:
+ self.label += ", +inf it/s"
+ else:
+ self.label += f", {1 / self.average_time_per_item:.2f} it/s"
+
+ self.label += "]" if self.comment is None or len(self.comment) == 0 else f", {self.comment}]"
+ self.display()
+
+ def display(self):
+ self.html_code = html_progress_bar(self.value, self.total, self.prefix, self.label, self.width)
+ if self.parent is not None:
+ # If this is a child bar, the parent will take care of the display.
+ self.parent.display()
+ return
+ if self.output is None:
+ self.output = disp.display(disp.HTML(self.html_code), display_id=True)
+ else:
+ self.output.update(disp.HTML(self.html_code))
+
+ def close(self):
+ "Closes the progress bar."
+ if self.parent is None and self.output is not None:
+ self.output.update(disp.HTML(""))
+
+
+class NotebookTrainingTracker(NotebookProgressBar):
+ """
+ An object tracking the updates of an ongoing training with progress bars and a nice table reporting metrics.
+
+ Args:
+ num_steps (`int`): The number of steps during training. column_names (`List[str]`, *optional*):
+ The list of column names for the metrics table (will be inferred from the first call to
+ [`~utils.notebook.NotebookTrainingTracker.write_line`] if not set).
+ """
+
+ def __init__(self, num_steps, column_names=None):
+ super().__init__(num_steps)
+ self.inner_table = None if column_names is None else [column_names]
+ self.child_bar = None
+
+ def display(self):
+ self.html_code = html_progress_bar(self.value, self.total, self.prefix, self.label, self.width)
+ if self.inner_table is not None:
+ self.html_code += text_to_html_table(self.inner_table)
+ if self.child_bar is not None:
+ self.html_code += self.child_bar.html_code
+ if self.output is None:
+ self.output = disp.display(disp.HTML(self.html_code), display_id=True)
+ else:
+ self.output.update(disp.HTML(self.html_code))
+
+ def write_line(self, values):
+ """
+ Write the values in the inner table.
+
+ Args:
+ values (`Dict[str, float]`): The values to display.
+ """
+ if self.inner_table is None:
+ self.inner_table = [list(values.keys()), list(values.values())]
+ else:
+ columns = self.inner_table[0]
+ for key in values.keys():
+ if key not in columns:
+ columns.append(key)
+ self.inner_table[0] = columns
+ if len(self.inner_table) > 1:
+ last_values = self.inner_table[-1]
+ first_column = self.inner_table[0][0]
+ if last_values[0] != values[first_column]:
+ # write new line
+ self.inner_table.append([values[c] if c in values else "No Log" for c in columns])
+ else:
+ # update last line
+ new_values = values
+ for c in columns:
+ if c not in new_values.keys():
+ new_values[c] = last_values[columns.index(c)]
+ self.inner_table[-1] = [new_values[c] for c in columns]
+ else:
+ self.inner_table.append([values[c] for c in columns])
+
+ def add_child(self, total, prefix=None, width=300):
+ """
+ Add a child progress bar displayed under the table of metrics. The child progress bar is returned (so it can be
+ easily updated).
+
+ Args:
+ total (`int`): The number of iterations for the child progress bar.
+ prefix (`str`, *optional*): A prefix to write on the left of the progress bar.
+ width (`int`, *optional*, defaults to 300): The width (in pixels) of the progress bar.
+ """
+ self.child_bar = NotebookProgressBar(total, prefix=prefix, parent=self, width=width)
+ return self.child_bar
+
+ def remove_child(self):
+ """
+ Closes the child progress bar.
+ """
+ self.child_bar = None
+ self.display()
+
+
+class NotebookProgressCallback(TrainerCallback):
+ """
+ A [`TrainerCallback`] that displays the progress of training or evaluation, optimized for Jupyter Notebooks or
+ Google colab.
+ """
+
+ def __init__(self):
+ self.training_tracker = None
+ self.prediction_bar = None
+ self._force_next_update = False
+
+ def on_train_begin(self, args, state, control, **kwargs):
+ self.first_column = "Epoch" if args.eval_strategy == IntervalStrategy.EPOCH else "Step"
+ self.training_loss = 0
+ self.last_log = 0
+ column_names = [self.first_column] + ["Training Loss"]
+ if args.eval_strategy != IntervalStrategy.NO:
+ column_names.append("Validation Loss")
+ self.training_tracker = NotebookTrainingTracker(state.max_steps, column_names)
+
+ def on_step_end(self, args, state, control, **kwargs):
+ epoch = int(state.epoch) if int(state.epoch) == state.epoch else f"{state.epoch:.2f}"
+ self.training_tracker.update(
+ state.global_step + 1,
+ comment=f"Epoch {epoch}/{state.num_train_epochs}",
+ force_update=self._force_next_update,
+ )
+ self._force_next_update = False
+
+ def on_prediction_step(self, args, state, control, eval_dataloader=None, **kwargs):
+ if not has_length(eval_dataloader):
+ return
+ if self.prediction_bar is None:
+ if self.training_tracker is not None:
+ self.prediction_bar = self.training_tracker.add_child(len(eval_dataloader))
+ else:
+ self.prediction_bar = NotebookProgressBar(len(eval_dataloader))
+ self.prediction_bar.update(1)
+ else:
+ self.prediction_bar.update(self.prediction_bar.value + 1)
+
+ def on_predict(self, args, state, control, **kwargs):
+ if self.prediction_bar is not None:
+ self.prediction_bar.close()
+ self.prediction_bar = None
+
+ def on_log(self, args, state, control, logs=None, **kwargs):
+ # Only for when there is no evaluation
+ if args.eval_strategy == IntervalStrategy.NO and "loss" in logs:
+ values = {"Training Loss": logs["loss"]}
+ # First column is necessarily Step sine we're not in epoch eval strategy
+ values["Step"] = state.global_step
+ self.training_tracker.write_line(values)
+
+ def on_evaluate(self, args, state, control, metrics=None, **kwargs):
+ if self.training_tracker is not None:
+ values = {"Training Loss": "No log", "Validation Loss": "No log"}
+ for log in reversed(state.log_history):
+ if "loss" in log:
+ values["Training Loss"] = log["loss"]
+ break
+
+ if self.first_column == "Epoch":
+ values["Epoch"] = int(state.epoch)
+ else:
+ values["Step"] = state.global_step
+ metric_key_prefix = "eval"
+ for k in metrics:
+ if k.endswith("_loss"):
+ metric_key_prefix = re.sub(r"\_loss$", "", k)
+ _ = metrics.pop("total_flos", None)
+ _ = metrics.pop("epoch", None)
+ _ = metrics.pop(f"{metric_key_prefix}_runtime", None)
+ _ = metrics.pop(f"{metric_key_prefix}_samples_per_second", None)
+ _ = metrics.pop(f"{metric_key_prefix}_steps_per_second", None)
+ _ = metrics.pop(f"{metric_key_prefix}_jit_compilation_time", None)
+ for k, v in metrics.items():
+ splits = k.split("_")
+ name = " ".join([part.capitalize() for part in splits[1:]])
+ if name == "Loss":
+ # Single dataset
+ name = "Validation Loss"
+ values[name] = v
+ self.training_tracker.write_line(values)
+ self.training_tracker.remove_child()
+ self.prediction_bar = None
+ # Evaluation takes a long time so we should force the next update.
+ self._force_next_update = True
+
+ def on_train_end(self, args, state, control, **kwargs):
+ self.training_tracker.update(
+ state.global_step,
+ comment=f"Epoch {int(state.epoch)}/{state.num_train_epochs}",
+ force_update=True,
+ )
+ self.training_tracker = None
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/npu_flash_attention.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/npu_flash_attention.py
new file mode 100644
index 0000000000000000000000000000000000000000..fa26379126fc09f9f56972741292c052a4895371
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/npu_flash_attention.py
@@ -0,0 +1,233 @@
+# 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 os
+
+import torch
+import torch.nn.functional as F
+
+from ..utils.import_utils import is_torch_npu_available
+
+
+if is_torch_npu_available():
+ import torch_npu
+ from einops import rearrange, repeat
+
+
+# FlashAttention2 is supported on Ascend NPU with down-right aligned causal mask by default.
+# Set environment variable `NPU_FA2_SPARSE_MODE` to 2 when using top-left aligned causal mask.
+TOP_LEFT_ALIGNED_CAUSAL_MASK_MODE = 2
+DOWN_RIGHT_ALIGNED_CAUSAL_MASK_MODE = 3
+
+SPARSE_MODE = int(os.getenv("NPU_FA2_SPARSE_MODE", default=DOWN_RIGHT_ALIGNED_CAUSAL_MASK_MODE))
+if SPARSE_MODE not in [TOP_LEFT_ALIGNED_CAUSAL_MASK_MODE, DOWN_RIGHT_ALIGNED_CAUSAL_MASK_MODE]:
+ raise ValueError(
+ "Environment variable `NPU_FA2_SPARSE_MODE` can only be set as 2 (top-left aligned causal mask) "
+ "or 3 (down-right aligned causal mask)."
+ )
+
+
+def is_npu_fa2_top_left_aligned_causal_mask():
+ return SPARSE_MODE == TOP_LEFT_ALIGNED_CAUSAL_MASK_MODE if is_torch_npu_available() else False
+
+
+# Copied from https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/bert_padding.py
+class IndexFirstAxis(torch.autograd.Function):
+ @staticmethod
+ def forward(ctx, input, indices):
+ ctx.save_for_backward(indices)
+ assert input.ndim >= 2
+ ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:]
+ second_dim = other_shape.numel()
+ # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing.
+ # return input[indices]
+ return torch.gather(
+ rearrange(input, "b ... -> b (...)"), 0, repeat(indices, "z -> z d", d=second_dim)
+ ).reshape(-1, *other_shape)
+
+ @staticmethod
+ def backward(ctx, grad_output):
+ (indices,) = ctx.saved_tensors
+ assert grad_output.ndim >= 2
+ other_shape = grad_output.shape[1:]
+ grad_output = rearrange(grad_output, "b ... -> b (...)")
+ grad_input = torch.zeros(
+ [ctx.first_axis_dim, grad_output.shape[1]],
+ device=grad_output.device,
+ dtype=grad_output.dtype,
+ )
+ # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing.
+ # grad_input[indices] = grad_output
+ grad_input.scatter_(0, repeat(indices, "z -> z d", d=grad_output.shape[1]), grad_output)
+ return grad_input.reshape(ctx.first_axis_dim, *other_shape), None
+
+
+index_first_axis = IndexFirstAxis.apply
+
+
+# Copied from https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/bert_padding.py
+class IndexPutFirstAxis(torch.autograd.Function):
+ @staticmethod
+ def forward(ctx, values, indices, first_axis_dim):
+ ctx.save_for_backward(indices)
+ assert indices.ndim == 1
+ assert values.ndim >= 2
+ output = torch.zeros(first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype)
+ # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing.
+ output[indices] = values
+ # output.scatter_(0, repeat(indices, 'z -> z d', d=values.shape[1]), values)
+ return output
+
+ @staticmethod
+ def backward(ctx, grad_output):
+ (indices,) = ctx.saved_tensors
+ # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing.
+ grad_values = grad_output[indices]
+ # grad_values = torch.gather(grad_output, 0, repeat(indices, 'z -> z d', d=grad_output.shape[1]))
+ return grad_values, None, None
+
+
+index_put_first_axis = IndexPutFirstAxis.apply
+
+
+# Copied from https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/bert_padding.py
+def pad_input(hidden_states, indices, batch, seqlen):
+ """
+ Arguments:
+ hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask.
+ indices: (total_nnz), the indices that represent the non-masked tokens of the original padded input sequence.
+ batch: int, batch size for the padded sequence.
+ seqlen: int, maximum sequence length for the padded sequence.
+ Return:
+ hidden_states: (batch, seqlen, ...)
+ """
+ # dim = hidden_states.shape[-1]
+ # output = torch.zeros((batch * seqlen), dim, device=hidden_states.device, dtype=hidden_states.dtype)
+ # output[indices] = hidden_states
+ output = index_put_first_axis(hidden_states, indices, batch * seqlen)
+ return rearrange(output, "(b s) ... -> b s ...", b=batch)
+
+
+# Copied from https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/bert_padding.py
+def unpad_input(hidden_states, attention_mask, unused_mask=None):
+ """
+ Arguments:
+ hidden_states: (batch, seqlen, ...)
+ attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid.
+ unused_mask: (batch, seqlen), bool / int, 1 means the element is allocated but unused.
+ Return:
+ hidden_states: (total_nnz, ...), where total_nnz = number of tokens selected in attention_mask + unused_mask.
+ indices: (total_nnz), the indices of masked tokens from the flattened input sequence.
+ cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states.
+ max_seqlen_in_batch: int
+ seqused: (batch), returns the number of tokens selected in attention_mask + unused_mask.
+ """
+ all_masks = (attention_mask + unused_mask) if unused_mask is not None else attention_mask
+ seqlens_in_batch = all_masks.sum(dim=-1, dtype=torch.int32)
+ used_seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
+ indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten()
+ max_seqlen_in_batch = seqlens_in_batch.max().item()
+ cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
+ # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the
+ # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim
+ # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to
+ # index with integer indices. Moreover, torch's index is a bit slower than it needs to be,
+ # so we write custom forward and backward to make it a bit faster.
+ return (
+ index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices),
+ indices,
+ cu_seqlens,
+ max_seqlen_in_batch,
+ used_seqlens_in_batch,
+ )
+
+
+def npu_flash_attn_func(
+ q,
+ k,
+ v,
+ dropout_p=0.0,
+ softmax_scale=None,
+ causal=False,
+ **kwargs,
+):
+ keep_prob = 1.0 - dropout_p
+
+ if not causal:
+ head_num = q.shape[2]
+ output = torch_npu.npu_fusion_attention(q, k, v, head_num, "BSND", keep_prob=keep_prob, scale=softmax_scale)[0]
+ else:
+ attn_mask_npu = torch.triu(torch.ones([2048, 2048]), diagonal=1).bool().to(q.device)
+ head_num = q.shape[2]
+ output = torch_npu.npu_fusion_attention(
+ q,
+ k,
+ v,
+ head_num,
+ "BSND",
+ keep_prob=keep_prob,
+ scale=softmax_scale,
+ atten_mask=attn_mask_npu,
+ sparse_mode=SPARSE_MODE,
+ )[0]
+
+ return output
+
+
+def npu_flash_attn_varlen_func(
+ q,
+ k,
+ v,
+ cu_seqlens_q,
+ cu_seqlens_k,
+ dropout_p=0.0,
+ softmax_scale=None,
+ causal=False,
+ **kwargs,
+):
+ keep_prob = 1.0 - dropout_p
+
+ if not causal:
+ head_num = q.shape[1]
+ output = torch_npu.npu_fusion_attention(
+ q,
+ k,
+ v,
+ head_num,
+ pse=None,
+ atten_mask=None,
+ scale=softmax_scale,
+ keep_prob=keep_prob,
+ input_layout="TND",
+ actual_seq_qlen=tuple(cu_seqlens_q[1:].cpu().numpy().tolist()),
+ actual_seq_kvlen=tuple(cu_seqlens_k[1:].cpu().numpy().tolist()),
+ )[0]
+ else:
+ attn_mask_npu = torch.triu(torch.ones([2048, 2048]), diagonal=1).bool().to(q.device)
+ head_num = q.shape[1]
+ output = torch_npu.npu_fusion_attention(
+ q,
+ k,
+ v,
+ head_num,
+ pse=None,
+ padding_mask=None,
+ atten_mask=attn_mask_npu,
+ scale=softmax_scale,
+ keep_prob=keep_prob,
+ input_layout="TND",
+ actual_seq_qlen=tuple(cu_seqlens_q[1:].cpu().numpy().tolist()),
+ actual_seq_kvlen=tuple(cu_seqlens_k[1:].cpu().numpy().tolist()),
+ sparse_mode=SPARSE_MODE,
+ )[0]
+
+ return output
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/optimization.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/optimization.py
new file mode 100644
index 0000000000000000000000000000000000000000..ba9d9920c1c99a50aa16d011bad68d012cba7eb7
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/optimization.py
@@ -0,0 +1,892 @@
+# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. 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.
+"""PyTorch optimization for BERT model."""
+
+import math
+import warnings
+from functools import partial
+from typing import Optional, Union
+
+import torch
+from torch.optim import Optimizer
+from torch.optim.lr_scheduler import LambdaLR, ReduceLROnPlateau
+
+from .trainer_pt_utils import LayerWiseDummyOptimizer, LayerWiseDummyScheduler
+from .trainer_utils import SchedulerType
+from .utils import logging
+from .utils.versions import require_version
+
+
+logger = logging.get_logger(__name__)
+
+
+def _get_constant_lambda(_=None):
+ return 1
+
+
+def get_constant_schedule(optimizer: Optimizer, last_epoch: int = -1):
+ """
+ Create a schedule with a constant learning rate, using the learning rate set in optimizer.
+
+ Args:
+ optimizer ([`~torch.optim.Optimizer`]):
+ The optimizer for which to schedule the learning rate.
+ last_epoch (`int`, *optional*, defaults to -1):
+ The index of the last epoch when resuming training.
+
+ Return:
+ `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
+ """
+
+ return LambdaLR(optimizer, _get_constant_lambda, last_epoch=last_epoch)
+
+
+def get_reduce_on_plateau_schedule(optimizer: Optimizer, **kwargs):
+ """
+ Create a schedule with a constant learning rate that decreases when a metric has stopped improving.
+
+ Args:
+ optimizer ([`~torch.optim.Optimizer`]):
+ The optimizer for which to schedule the learning rate.
+ kwargs (`dict`, *optional*):
+ Extra parameters to be passed to the scheduler. See `torch.optim.lr_scheduler.ReduceLROnPlateau`
+ for possible parameters.
+
+ Return:
+ `torch.optim.lr_scheduler.ReduceLROnPlateau` with the appropriate schedule.
+ """
+
+ return ReduceLROnPlateau(optimizer, **kwargs)
+
+
+def _get_constant_schedule_with_warmup_lr_lambda(current_step: int, *, num_warmup_steps: int):
+ if current_step < num_warmup_steps:
+ return float(current_step) / float(max(1.0, num_warmup_steps))
+ return 1.0
+
+
+def get_constant_schedule_with_warmup(optimizer: Optimizer, num_warmup_steps: int, last_epoch: int = -1):
+ """
+ Create a schedule with a constant learning rate preceded by a warmup period during which the learning rate
+ increases linearly between 0 and the initial lr set in the optimizer.
+
+ Args:
+ optimizer ([`~torch.optim.Optimizer`]):
+ The optimizer for which to schedule the learning rate.
+ num_warmup_steps (`int`):
+ The number of steps for the warmup phase.
+ last_epoch (`int`, *optional*, defaults to -1):
+ The index of the last epoch when resuming training.
+
+ Return:
+ `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
+ """
+
+ lr_lambda = partial(_get_constant_schedule_with_warmup_lr_lambda, num_warmup_steps=num_warmup_steps)
+ return LambdaLR(optimizer, lr_lambda, last_epoch=last_epoch)
+
+
+def _get_linear_schedule_with_warmup_lr_lambda(current_step: int, *, num_warmup_steps: int, num_training_steps: int):
+ if current_step < num_warmup_steps:
+ return float(current_step) / float(max(1, num_warmup_steps))
+ return max(0.0, float(num_training_steps - current_step) / float(max(1, num_training_steps - num_warmup_steps)))
+
+
+def get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps, last_epoch=-1):
+ """
+ Create a schedule with a learning rate that decreases linearly from the initial lr set in the optimizer to 0, after
+ a warmup period during which it increases linearly from 0 to the initial lr set in the optimizer.
+
+ Args:
+ optimizer ([`~torch.optim.Optimizer`]):
+ The optimizer for which to schedule the learning rate.
+ num_warmup_steps (`int`):
+ The number of steps for the warmup phase.
+ num_training_steps (`int`):
+ The total number of training steps.
+ last_epoch (`int`, *optional*, defaults to -1):
+ The index of the last epoch when resuming training.
+
+ Return:
+ `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
+ """
+
+ lr_lambda = partial(
+ _get_linear_schedule_with_warmup_lr_lambda,
+ num_warmup_steps=num_warmup_steps,
+ num_training_steps=num_training_steps,
+ )
+ return LambdaLR(optimizer, lr_lambda, last_epoch)
+
+
+def _get_cosine_schedule_with_warmup_lr_lambda(
+ current_step: int, *, num_warmup_steps: int, num_training_steps: int, num_cycles: float
+):
+ if current_step < num_warmup_steps:
+ return float(current_step) / float(max(1, num_warmup_steps))
+ progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
+ return max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress)))
+
+
+def get_cosine_schedule_with_warmup(
+ optimizer: Optimizer, num_warmup_steps: int, num_training_steps: int, num_cycles: float = 0.5, last_epoch: int = -1
+):
+ """
+ Create a schedule with a learning rate that decreases following the values of the cosine function between the
+ initial lr set in the optimizer to 0, after a warmup period during which it increases linearly between 0 and the
+ initial lr set in the optimizer.
+
+ Args:
+ optimizer ([`~torch.optim.Optimizer`]):
+ The optimizer for which to schedule the learning rate.
+ num_warmup_steps (`int`):
+ The number of steps for the warmup phase.
+ num_training_steps (`int`):
+ The total number of training steps.
+ num_cycles (`float`, *optional*, defaults to 0.5):
+ The number of waves in the cosine schedule (the defaults is to just decrease from the max value to 0
+ following a half-cosine).
+ last_epoch (`int`, *optional*, defaults to -1):
+ The index of the last epoch when resuming training.
+
+ Return:
+ `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
+ """
+
+ lr_lambda = partial(
+ _get_cosine_schedule_with_warmup_lr_lambda,
+ num_warmup_steps=num_warmup_steps,
+ num_training_steps=num_training_steps,
+ num_cycles=num_cycles,
+ )
+ return LambdaLR(optimizer, lr_lambda, last_epoch)
+
+
+def _get_cosine_with_hard_restarts_schedule_with_warmup_lr_lambda(
+ current_step: int, *, num_warmup_steps: int, num_training_steps: int, num_cycles: int
+):
+ if current_step < num_warmup_steps:
+ return float(current_step) / float(max(1, num_warmup_steps))
+ progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
+ if progress >= 1.0:
+ return 0.0
+ return max(0.0, 0.5 * (1.0 + math.cos(math.pi * ((float(num_cycles) * progress) % 1.0))))
+
+
+def get_cosine_with_hard_restarts_schedule_with_warmup(
+ optimizer: Optimizer, num_warmup_steps: int, num_training_steps: int, num_cycles: int = 1, last_epoch: int = -1
+):
+ """
+ Create a schedule with a learning rate that decreases following the values of the cosine function between the
+ initial lr set in the optimizer to 0, with several hard restarts, after a warmup period during which it increases
+ linearly between 0 and the initial lr set in the optimizer.
+
+ Args:
+ optimizer ([`~torch.optim.Optimizer`]):
+ The optimizer for which to schedule the learning rate.
+ num_warmup_steps (`int`):
+ The number of steps for the warmup phase.
+ num_training_steps (`int`):
+ The total number of training steps.
+ num_cycles (`int`, *optional*, defaults to 1):
+ The number of hard restarts to use.
+ last_epoch (`int`, *optional*, defaults to -1):
+ The index of the last epoch when resuming training.
+
+ Return:
+ `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
+ """
+
+ lr_lambda = partial(
+ _get_cosine_with_hard_restarts_schedule_with_warmup_lr_lambda,
+ num_warmup_steps=num_warmup_steps,
+ num_training_steps=num_training_steps,
+ num_cycles=num_cycles,
+ )
+ return LambdaLR(optimizer, lr_lambda, last_epoch)
+
+
+def _get_polynomial_decay_schedule_with_warmup_lr_lambda(
+ current_step: int,
+ *,
+ num_warmup_steps: int,
+ num_training_steps: int,
+ lr_end: float,
+ power: float,
+ lr_init: int,
+):
+ if current_step < num_warmup_steps:
+ return float(current_step) / float(max(1, num_warmup_steps))
+ elif current_step > num_training_steps:
+ return lr_end / lr_init # as LambdaLR multiplies by lr_init
+ else:
+ lr_range = lr_init - lr_end
+ decay_steps = num_training_steps - num_warmup_steps
+ pct_remaining = 1 - (current_step - num_warmup_steps) / decay_steps
+ decay = lr_range * pct_remaining**power + lr_end
+ return decay / lr_init # as LambdaLR multiplies by lr_init
+
+
+def get_polynomial_decay_schedule_with_warmup(
+ optimizer, num_warmup_steps, num_training_steps, lr_end=1e-7, power=1.0, last_epoch=-1
+):
+ """
+ Create a schedule with a learning rate that decreases as a polynomial decay from the initial lr set in the
+ optimizer to end lr defined by *lr_end*, after a warmup period during which it increases linearly from 0 to the
+ initial lr set in the optimizer.
+
+ Args:
+ optimizer ([`~torch.optim.Optimizer`]):
+ The optimizer for which to schedule the learning rate.
+ num_warmup_steps (`int`):
+ The number of steps for the warmup phase.
+ num_training_steps (`int`):
+ The total number of training steps.
+ lr_end (`float`, *optional*, defaults to 1e-7):
+ The end LR.
+ power (`float`, *optional*, defaults to 1.0):
+ Power factor.
+ last_epoch (`int`, *optional*, defaults to -1):
+ The index of the last epoch when resuming training.
+
+ Note: *power* defaults to 1.0 as in the fairseq implementation, which in turn is based on the original BERT
+ implementation at
+ https://github.com/google-research/bert/blob/f39e881b169b9d53bea03d2d341b31707a6c052b/optimization.py#L37
+
+ Return:
+ `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
+
+ """
+
+ lr_init = optimizer.defaults["lr"]
+ if not (lr_init > lr_end):
+ raise ValueError(f"lr_end ({lr_end}) must be smaller than initial lr ({lr_init})")
+
+ lr_lambda = partial(
+ _get_polynomial_decay_schedule_with_warmup_lr_lambda,
+ num_warmup_steps=num_warmup_steps,
+ num_training_steps=num_training_steps,
+ lr_end=lr_end,
+ power=power,
+ lr_init=lr_init,
+ )
+ return LambdaLR(optimizer, lr_lambda, last_epoch)
+
+
+def _get_inverse_sqrt_schedule_lr_lambda(current_step: int, *, num_warmup_steps: int, timescale: Optional[int] = None):
+ if current_step < num_warmup_steps:
+ return float(current_step) / float(max(1, num_warmup_steps))
+ shift = timescale - num_warmup_steps
+ decay = 1.0 / math.sqrt((current_step + shift) / timescale)
+ return decay
+
+
+def get_inverse_sqrt_schedule(
+ optimizer: Optimizer, num_warmup_steps: int, timescale: Optional[int] = None, last_epoch: int = -1
+):
+ """
+ Create a schedule with an inverse square-root learning rate, from the initial lr set in the optimizer, after a
+ warmup period which increases lr linearly from 0 to the initial lr set in the optimizer.
+
+ Args:
+ optimizer ([`~torch.optim.Optimizer`]):
+ The optimizer for which to schedule the learning rate.
+ num_warmup_steps (`int`):
+ The number of steps for the warmup phase.
+ timescale (`int`, *optional*, defaults to `num_warmup_steps`):
+ Time scale.
+ last_epoch (`int`, *optional*, defaults to -1):
+ The index of the last epoch when resuming training.
+
+ Return:
+ `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
+ """
+ # Note: this implementation is adapted from
+ # https://github.com/google-research/big_vision/blob/f071ce68852d56099437004fd70057597a95f6ef/big_vision/utils.py#L930
+
+ if timescale is None:
+ timescale = num_warmup_steps or 10_000
+
+ lr_lambda = partial(_get_inverse_sqrt_schedule_lr_lambda, num_warmup_steps=num_warmup_steps, timescale=timescale)
+ return LambdaLR(optimizer, lr_lambda, last_epoch=last_epoch)
+
+
+def _get_cosine_schedule_with_warmup_lr_lambda(
+ current_step: int, *, num_warmup_steps: int, num_training_steps: int, num_cycles: float, min_lr_rate: float = 0.0
+):
+ if current_step < num_warmup_steps:
+ return float(current_step) / float(max(1, num_warmup_steps))
+ progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
+ factor = 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress))
+ factor = factor * (1 - min_lr_rate) + min_lr_rate
+ return max(0, factor)
+
+
+def get_cosine_with_min_lr_schedule_with_warmup(
+ optimizer: Optimizer,
+ num_warmup_steps: int,
+ num_training_steps: int,
+ num_cycles: float = 0.5,
+ last_epoch: int = -1,
+ min_lr: Optional[float] = None,
+ min_lr_rate: Optional[float] = None,
+):
+ """
+ Create a schedule with a learning rate that decreases following the values of the cosine function between the
+ initial lr set in the optimizer to min_lr, after a warmup period during which it increases linearly between 0 and the
+ initial lr set in the optimizer.
+
+ Args:
+ optimizer ([`~torch.optim.Optimizer`]):
+ The optimizer for which to schedule the learning rate.
+ num_warmup_steps (`int`):
+ The number of steps for the warmup phase.
+ num_training_steps (`int`):
+ The total number of training steps.
+ num_cycles (`float`, *optional*, defaults to 0.5):
+ The number of waves in the cosine schedule (the defaults is to just decrease from the max value to 0
+ following a half-cosine).
+ last_epoch (`int`, *optional*, defaults to -1):
+ The index of the last epoch when resuming training.
+ min_lr (`float`, *optional*):
+ The minimum learning rate to reach after the cosine schedule.
+ min_lr_rate (`float`, *optional*):
+ The minimum learning rate as a ratio of the initial learning rate. If set, `min_lr` should not be set.
+
+ Return:
+ `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
+ """
+
+ if min_lr is not None and min_lr_rate is not None:
+ raise ValueError("Only one of min_lr or min_lr_rate should be set")
+ elif min_lr is not None:
+ min_lr_rate = min_lr / optimizer.defaults["lr"]
+ elif min_lr_rate is None:
+ raise ValueError("One of min_lr or min_lr_rate should be set through the `lr_scheduler_kwargs`")
+
+ lr_lambda = partial(
+ _get_cosine_schedule_with_warmup_lr_lambda,
+ num_warmup_steps=num_warmup_steps,
+ num_training_steps=num_training_steps,
+ num_cycles=num_cycles,
+ min_lr_rate=min_lr_rate,
+ )
+ return LambdaLR(optimizer, lr_lambda, last_epoch)
+
+
+def _get_wsd_scheduler_lambda(
+ current_step: int,
+ *,
+ num_warmup_steps: int,
+ num_stable_steps: int,
+ num_decay_steps: int,
+ warmup_type: str,
+ decay_type: str,
+ min_lr_ratio: float,
+ num_cycles: float,
+):
+ if current_step < num_warmup_steps:
+ progress = float(current_step) / float(max(1, num_warmup_steps))
+ if warmup_type == "linear":
+ factor = progress
+ elif warmup_type == "cosine":
+ factor = 0.5 * (1.0 - math.cos(math.pi * progress))
+ elif warmup_type == "1-sqrt":
+ factor = 1.0 - math.sqrt(1.0 - progress)
+ factor = factor * (1.0 - min_lr_ratio) + min_lr_ratio
+ return max(0.0, factor)
+
+ if current_step < num_warmup_steps + num_stable_steps:
+ return 1.0
+
+ if current_step < num_warmup_steps + num_stable_steps + num_decay_steps:
+ progress = float(current_step - num_warmup_steps - num_stable_steps) / float(max(1, num_decay_steps))
+ if decay_type == "linear":
+ factor = 1.0 - progress
+ elif decay_type == "cosine":
+ factor = 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress))
+ elif decay_type == "1-sqrt":
+ factor = 1.0 - math.sqrt(progress)
+ factor = factor * (1.0 - min_lr_ratio) + min_lr_ratio
+ return max(0.0, factor)
+ return min_lr_ratio
+
+
+def get_wsd_schedule(
+ optimizer: Optimizer,
+ num_warmup_steps: int,
+ num_decay_steps: int,
+ num_training_steps: Optional[int] = None,
+ num_stable_steps: Optional[int] = None,
+ warmup_type: str = "linear",
+ decay_type: str = "cosine",
+ min_lr_ratio: float = 0,
+ num_cycles: float = 0.5,
+ last_epoch: int = -1,
+):
+ """
+ Create a schedule with a learning rate that has three stages:
+ 1. warmup: increase from min_lr_ratio times the initial learning rate to the initial learning rate following a warmup_type.
+ 2. stable: constant learning rate.
+ 3. decay: decrease from the initial learning rate to min_lr_ratio times the initial learning rate following a decay_type.
+
+ Args:
+ optimizer ([`~torch.optim.Optimizer`]):
+ The optimizer for which to schedule the learning rate.
+ num_warmup_steps (`int`):
+ The number of steps for the warmup phase.
+ num_decay_steps (`int`):
+ The number of steps for the decay phase.
+ num_training_steps (`int`, *optional*):
+ The total number of training steps. This is the sum of the warmup, stable and decay steps. If `num_stable_steps` is not provided, the stable phase will be `num_training_steps - num_warmup_steps - num_decay_steps`.
+ num_stable_steps (`int`, *optional*):
+ The number of steps for the stable phase. Please ensure that `num_warmup_steps + num_stable_steps + num_decay_steps` equals `num_training_steps`, otherwise the other steps will default to the minimum learning rate.
+ warmup_type (`str`, *optional*, defaults to "linear"):
+ The type of warmup to use. Can be 'linear', 'cosine' or '1-sqrt'.
+ decay_type (`str`, *optional*, defaults to "cosine"):
+ The type of decay to use. Can be 'linear', 'cosine' or '1-sqrt'.
+ min_lr_ratio (`float`, *optional*, defaults to 0):
+ The minimum learning rate as a ratio of the initial learning rate.
+ num_cycles (`float`, *optional*, defaults to 0.5):
+ The number of waves in the cosine schedule (the defaults is to just decrease from the max value to 0
+ following a half-cosine).
+ last_epoch (`int`, *optional*, defaults to -1):
+ The index of the last epoch when resuming training.
+
+ Return:
+ `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
+ """
+
+ if num_training_steps is None and num_stable_steps is None:
+ raise ValueError("Either num_training_steps or num_stable_steps must be specified.")
+
+ if num_training_steps is not None and num_stable_steps is not None:
+ warnings.warn("Both num_training_steps and num_stable_steps are specified. num_stable_steps will be used.")
+
+ if warmup_type not in ["linear", "cosine", "1-sqrt"]:
+ raise ValueError(f"Unknown warmup type: {warmup_type}, expected 'linear', 'cosine' or '1-sqrt'")
+
+ if decay_type not in ["linear", "cosine", "1-sqrt"]:
+ raise ValueError(f"Unknown decay type: {decay_type}, expected 'linear', 'cosine' or '1-sqrt'")
+
+ if num_stable_steps is None:
+ num_stable_steps = num_training_steps - num_warmup_steps - num_decay_steps
+
+ lr_lambda = partial(
+ _get_wsd_scheduler_lambda,
+ num_warmup_steps=num_warmup_steps,
+ num_stable_steps=num_stable_steps,
+ num_decay_steps=num_decay_steps,
+ warmup_type=warmup_type,
+ decay_type=decay_type,
+ min_lr_ratio=min_lr_ratio,
+ num_cycles=num_cycles,
+ )
+ return LambdaLR(optimizer, lr_lambda, last_epoch)
+
+
+TYPE_TO_SCHEDULER_FUNCTION = {
+ SchedulerType.LINEAR: get_linear_schedule_with_warmup,
+ SchedulerType.COSINE: get_cosine_schedule_with_warmup,
+ SchedulerType.COSINE_WITH_RESTARTS: get_cosine_with_hard_restarts_schedule_with_warmup,
+ SchedulerType.POLYNOMIAL: get_polynomial_decay_schedule_with_warmup,
+ SchedulerType.CONSTANT: get_constant_schedule,
+ SchedulerType.CONSTANT_WITH_WARMUP: get_constant_schedule_with_warmup,
+ SchedulerType.INVERSE_SQRT: get_inverse_sqrt_schedule,
+ SchedulerType.REDUCE_ON_PLATEAU: get_reduce_on_plateau_schedule,
+ SchedulerType.COSINE_WITH_MIN_LR: get_cosine_with_min_lr_schedule_with_warmup,
+ SchedulerType.WARMUP_STABLE_DECAY: get_wsd_schedule,
+}
+
+
+def get_scheduler(
+ name: Union[str, SchedulerType],
+ optimizer: Optimizer,
+ num_warmup_steps: Optional[int] = None,
+ num_training_steps: Optional[int] = None,
+ scheduler_specific_kwargs: Optional[dict] = None,
+):
+ """
+ Unified API to get any scheduler from its name.
+
+ Args:
+ name (`str` or `SchedulerType`):
+ The name of the scheduler to use.
+ optimizer (`torch.optim.Optimizer`):
+ The optimizer that will be used during training.
+ num_warmup_steps (`int`, *optional*):
+ The number of warmup steps to do. This is not required by all schedulers (hence the argument being
+ optional), the function will raise an error if it's unset and the scheduler type requires it.
+ num_training_steps (`int``, *optional*):
+ The number of training steps to do. This is not required by all schedulers (hence the argument being
+ optional), the function will raise an error if it's unset and the scheduler type requires it.
+ scheduler_specific_kwargs (`dict`, *optional*):
+ Extra parameters for schedulers such as cosine with restarts. Mismatched scheduler types and scheduler
+ parameters will cause the scheduler function to raise a TypeError.
+ """
+ name = SchedulerType(name)
+ schedule_func = TYPE_TO_SCHEDULER_FUNCTION[name]
+
+ # If a `LayerWiseDummyOptimizer` is passed we extract the optimizer dict and
+ # recursively call `get_scheduler` to get the proper schedulers on each parameter
+ if optimizer is not None and isinstance(optimizer, LayerWiseDummyOptimizer):
+ optimizer_dict = optimizer.optimizer_dict
+ scheduler_dict = {}
+
+ for param in optimizer_dict.keys():
+ scheduler_dict[param] = get_scheduler(
+ name,
+ optimizer=optimizer_dict[param],
+ num_warmup_steps=num_warmup_steps,
+ num_training_steps=num_training_steps,
+ )
+
+ def scheduler_hook(param):
+ # Since the optimizer hook has been already attached we only need to
+ # attach the scheduler hook, the gradients have been zeroed here
+ scheduler_dict[param].step()
+
+ for param in optimizer_dict.keys():
+ if param.requires_grad:
+ param.register_post_accumulate_grad_hook(scheduler_hook)
+
+ return LayerWiseDummyScheduler(optimizer_dict=optimizer_dict, lr=optimizer.defaults["lr"])
+
+ if name == SchedulerType.CONSTANT:
+ return schedule_func(optimizer)
+
+ if scheduler_specific_kwargs is None:
+ scheduler_specific_kwargs = {}
+
+ if name == SchedulerType.REDUCE_ON_PLATEAU:
+ return schedule_func(optimizer, **scheduler_specific_kwargs)
+
+ # All other schedulers require `num_warmup_steps`
+ if num_warmup_steps is None:
+ raise ValueError(f"{name} requires `num_warmup_steps`, please provide that argument.")
+
+ if name == SchedulerType.CONSTANT_WITH_WARMUP:
+ return schedule_func(optimizer, num_warmup_steps=num_warmup_steps)
+
+ if name == SchedulerType.INVERSE_SQRT:
+ return schedule_func(optimizer, num_warmup_steps=num_warmup_steps)
+
+ # wsd scheduler requires either num_training_steps or num_stable_steps
+ if name == SchedulerType.WARMUP_STABLE_DECAY:
+ return schedule_func(
+ optimizer,
+ num_warmup_steps=num_warmup_steps,
+ num_training_steps=num_training_steps,
+ **scheduler_specific_kwargs,
+ )
+
+ # All other schedulers require `num_training_steps`
+ if num_training_steps is None:
+ raise ValueError(f"{name} requires `num_training_steps`, please provide that argument.")
+
+ return schedule_func(
+ optimizer,
+ num_warmup_steps=num_warmup_steps,
+ num_training_steps=num_training_steps,
+ **scheduler_specific_kwargs,
+ )
+
+
+class Adafactor(Optimizer):
+ """
+ AdaFactor pytorch implementation can be used as a drop in replacement for Adam original fairseq code:
+ https://github.com/pytorch/fairseq/blob/master/fairseq/optim/adafactor.py
+
+ Paper: *Adafactor: Adaptive Learning Rates with Sublinear Memory Cost* https://arxiv.org/abs/1804.04235 Note that
+ this optimizer internally adjusts the learning rate depending on the `scale_parameter`, `relative_step` and
+ `warmup_init` options. To use a manual (external) learning rate schedule you should set `scale_parameter=False` and
+ `relative_step=False`.
+
+ Arguments:
+ params (`Iterable[nn.parameter.Parameter]`):
+ Iterable of parameters to optimize or dictionaries defining parameter groups.
+ lr (`float`, *optional*):
+ The external learning rate.
+ eps (`Tuple[float, float]`, *optional*, defaults to `(1e-30, 0.001)`):
+ Regularization constants for square gradient and parameter scale respectively
+ clip_threshold (`float`, *optional*, defaults to 1.0):
+ Threshold of root mean square of final gradient update
+ decay_rate (`float`, *optional*, defaults to -0.8):
+ Coefficient used to compute running averages of square
+ beta1 (`float`, *optional*):
+ Coefficient used for computing running averages of gradient
+ weight_decay (`float`, *optional*, defaults to 0.0):
+ Weight decay (L2 penalty)
+ scale_parameter (`bool`, *optional*, defaults to `True`):
+ If True, learning rate is scaled by root mean square
+ relative_step (`bool`, *optional*, defaults to `True`):
+ If True, time-dependent learning rate is computed instead of external learning rate
+ warmup_init (`bool`, *optional*, defaults to `False`):
+ Time-dependent learning rate computation depends on whether warm-up initialization is being used
+
+ This implementation handles low-precision (FP16, bfloat) values, but we have not thoroughly tested.
+
+ Recommended T5 finetuning settings (https://discuss.huggingface.co/t/t5-finetuning-tips/684/3):
+
+ - Training without LR warmup or clip_threshold is not recommended.
+
+ - use scheduled LR warm-up to fixed LR
+ - use clip_threshold=1.0 (https://arxiv.org/abs/1804.04235)
+ - Disable relative updates
+ - Use scale_parameter=False
+ - Additional optimizer operations like gradient clipping should not be used alongside Adafactor
+
+ Example:
+
+ ```python
+ Adafactor(model.parameters(), scale_parameter=False, relative_step=False, warmup_init=False, lr=1e-3)
+ ```
+
+ Others reported the following combination to work well:
+
+ ```python
+ Adafactor(model.parameters(), scale_parameter=True, relative_step=True, warmup_init=True, lr=None)
+ ```
+
+ When using `lr=None` with [`Trainer`] you will most likely need to use [`~optimization.AdafactorSchedule`]
+ scheduler as following:
+
+ ```python
+ from transformers.optimization import Adafactor, AdafactorSchedule
+
+ optimizer = Adafactor(model.parameters(), scale_parameter=True, relative_step=True, warmup_init=True, lr=None)
+ lr_scheduler = AdafactorSchedule(optimizer)
+ trainer = Trainer(..., optimizers=(optimizer, lr_scheduler))
+ ```
+
+ Usage:
+
+ ```python
+ # replace AdamW with Adafactor
+ optimizer = Adafactor(
+ model.parameters(),
+ lr=1e-3,
+ eps=(1e-30, 1e-3),
+ clip_threshold=1.0,
+ decay_rate=-0.8,
+ beta1=None,
+ weight_decay=0.0,
+ relative_step=False,
+ scale_parameter=False,
+ warmup_init=False,
+ )
+ ```"""
+
+ def __init__(
+ self,
+ params,
+ lr=None,
+ eps=(1e-30, 1e-3),
+ clip_threshold=1.0,
+ decay_rate=-0.8,
+ beta1=None,
+ weight_decay=0.0,
+ scale_parameter=True,
+ relative_step=True,
+ warmup_init=False,
+ ):
+ require_version("torch>=1.5.0") # add_ with alpha
+ if lr is not None and relative_step:
+ raise ValueError("Cannot combine manual `lr` and `relative_step=True` options")
+ if warmup_init and not relative_step:
+ raise ValueError("`warmup_init=True` requires `relative_step=True`")
+
+ defaults = {
+ "lr": lr,
+ "eps": eps,
+ "clip_threshold": clip_threshold,
+ "decay_rate": decay_rate,
+ "beta1": beta1,
+ "weight_decay": weight_decay,
+ "scale_parameter": scale_parameter,
+ "relative_step": relative_step,
+ "warmup_init": warmup_init,
+ }
+ super().__init__(params, defaults)
+
+ @staticmethod
+ def _get_lr(param_group, param_state):
+ rel_step_sz = param_group["lr"]
+ if param_group["relative_step"]:
+ min_step = 1e-6 * param_state["step"] if param_group["warmup_init"] else 1e-2
+ rel_step_sz = min(min_step, 1.0 / math.sqrt(param_state["step"]))
+ param_scale = 1.0
+ if param_group["scale_parameter"]:
+ param_scale = max(param_group["eps"][1], param_state["RMS"])
+ return param_scale * rel_step_sz
+
+ @staticmethod
+ def _get_options(param_group, param_shape):
+ factored = len(param_shape) >= 2
+ use_first_moment = param_group["beta1"] is not None
+ return factored, use_first_moment
+
+ @staticmethod
+ def _rms(tensor):
+ return tensor.norm(2) / (tensor.numel() ** 0.5)
+
+ @staticmethod
+ def _approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col):
+ # copy from fairseq's adafactor implementation:
+ # https://github.com/huggingface/transformers/blob/8395f14de6068012787d83989c3627c3df6a252b/src/transformers/optimization.py#L505
+ r_factor = (exp_avg_sq_row / exp_avg_sq_row.mean(dim=-1, keepdim=True)).rsqrt_().unsqueeze(-1)
+ c_factor = exp_avg_sq_col.unsqueeze(-2).rsqrt()
+ return torch.mul(r_factor, c_factor)
+
+ @torch.no_grad()
+ def step(self, closure=None):
+ """
+ Performs a single optimization step
+
+ Arguments:
+ closure (callable, optional): A closure that reevaluates the model
+ and returns the loss.
+ """
+ loss = None
+ if closure is not None:
+ loss = closure()
+
+ for group in self.param_groups:
+ for p in group["params"]:
+ if p.grad is None:
+ continue
+ grad = p.grad
+ if grad.dtype in {torch.float16, torch.bfloat16}:
+ grad = grad.float()
+ if grad.is_sparse:
+ raise RuntimeError("Adafactor does not support sparse gradients.")
+
+ state = self.state[p]
+ grad_shape = grad.shape
+
+ factored, use_first_moment = self._get_options(group, grad_shape)
+ # State Initialization
+ if len(state) == 0:
+ state["step"] = 0
+
+ if use_first_moment:
+ # Exponential moving average of gradient values
+ state["exp_avg"] = torch.zeros_like(grad)
+ if factored:
+ state["exp_avg_sq_row"] = torch.zeros(grad_shape[:-1]).to(grad)
+ state["exp_avg_sq_col"] = torch.zeros(grad_shape[:-2] + grad_shape[-1:]).to(grad)
+ else:
+ state["exp_avg_sq"] = torch.zeros_like(grad)
+
+ state["RMS"] = 0
+ else:
+ if use_first_moment:
+ state["exp_avg"] = state["exp_avg"].to(grad)
+ if factored:
+ state["exp_avg_sq_row"] = state["exp_avg_sq_row"].to(grad)
+ state["exp_avg_sq_col"] = state["exp_avg_sq_col"].to(grad)
+ else:
+ state["exp_avg_sq"] = state["exp_avg_sq"].to(grad)
+
+ p_data_fp32 = p
+ if p.dtype in {torch.float16, torch.bfloat16}:
+ p_data_fp32 = p_data_fp32.float()
+
+ state["step"] += 1
+ state["RMS"] = self._rms(p_data_fp32)
+ lr = self._get_lr(group, state)
+
+ beta2t = 1.0 - math.pow(state["step"], group["decay_rate"])
+ update = (grad**2) + group["eps"][0]
+ if factored:
+ exp_avg_sq_row = state["exp_avg_sq_row"]
+ exp_avg_sq_col = state["exp_avg_sq_col"]
+
+ exp_avg_sq_row.mul_(beta2t).add_(update.mean(dim=-1), alpha=(1.0 - beta2t))
+ exp_avg_sq_col.mul_(beta2t).add_(update.mean(dim=-2), alpha=(1.0 - beta2t))
+
+ # Approximation of exponential moving average of square of gradient
+ update = self._approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col)
+ update.mul_(grad)
+ else:
+ exp_avg_sq = state["exp_avg_sq"]
+
+ exp_avg_sq.mul_(beta2t).add_(update, alpha=(1.0 - beta2t))
+ update = exp_avg_sq.rsqrt().mul_(grad)
+
+ update.div_((self._rms(update) / group["clip_threshold"]).clamp_(min=1.0))
+ update.mul_(lr)
+
+ if use_first_moment:
+ exp_avg = state["exp_avg"]
+ exp_avg.mul_(group["beta1"]).add_(update, alpha=(1 - group["beta1"]))
+ update = exp_avg
+
+ if group["weight_decay"] != 0:
+ p_data_fp32.add_(p_data_fp32, alpha=(-group["weight_decay"] * lr))
+
+ p_data_fp32.add_(-update)
+
+ if p.dtype in {torch.float16, torch.bfloat16}:
+ p.copy_(p_data_fp32)
+
+ return loss
+
+
+class AdafactorSchedule(LambdaLR):
+ """
+ Since [`~optimization.Adafactor`] performs its own scheduling, if the training loop relies on a scheduler (e.g.,
+ for logging), this class creates a proxy object that retrieves the current lr values from the optimizer.
+
+ It returns `initial_lr` during startup and the actual `lr` during stepping.
+ """
+
+ def __init__(self, optimizer, initial_lr=0.0):
+ def lr_lambda(_):
+ return initial_lr
+
+ for group in optimizer.param_groups:
+ group["initial_lr"] = initial_lr
+ super().__init__(optimizer, lr_lambda)
+ for group in optimizer.param_groups:
+ del group["initial_lr"]
+
+ def get_lr(self):
+ opt = self.optimizer
+ lrs = [
+ opt._get_lr(group, opt.state[group["params"][0]])
+ for group in opt.param_groups
+ if group["params"][0].grad is not None
+ ]
+ if len(lrs) == 0:
+ lrs = self.base_lrs # if called before stepping
+ return lrs
+
+
+def get_adafactor_schedule(optimizer, initial_lr=0.0):
+ """
+ Get a proxy schedule for [`~optimization.Adafactor`]
+
+ Args:
+ optimizer ([`~torch.optim.Optimizer`]):
+ The optimizer for which to schedule the learning rate.
+ initial_lr (`float`, *optional*, defaults to 0.0):
+ Initial lr
+
+ Return:
+ [`~optimization.Adafactor`] proxy schedule object.
+
+
+ """
+ return AdafactorSchedule(optimizer, initial_lr)
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/peft.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/peft.py
new file mode 100644
index 0000000000000000000000000000000000000000..6aa3b137b1914254c1b9ac49d5a8b413e29acb0b
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/peft.py
@@ -0,0 +1,593 @@
+# Copyright 2023 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.
+
+import importlib
+import inspect
+import warnings
+from typing import Any, Dict, List, Optional, Union
+
+from packaging import version
+
+from ..utils import (
+ check_peft_version,
+ find_adapter_config_file,
+ is_accelerate_available,
+ is_peft_available,
+ is_torch_available,
+ logging,
+)
+
+
+if is_torch_available():
+ import torch
+
+if is_accelerate_available():
+ from accelerate import dispatch_model
+ from accelerate.utils import get_balanced_memory, infer_auto_device_map
+
+# Minimum PEFT version supported for the integration
+MIN_PEFT_VERSION = "0.5.0"
+
+
+logger = logging.get_logger(__name__)
+
+
+class PeftAdapterMixin:
+ """
+ A class containing all functions for loading and using adapters weights that are supported in PEFT library. For
+ more details about adapters and injecting them on a transformer-based model, check out the documentation of PEFT
+ library: https://huggingface.co/docs/peft/index
+
+ Currently supported PEFT methods are all non-prefix tuning methods. Below is the list of supported PEFT methods
+ that anyone can load, train and run with this mixin class:
+ - Low Rank Adapters (LoRA): https://huggingface.co/docs/peft/conceptual_guides/lora
+ - IA3: https://huggingface.co/docs/peft/conceptual_guides/ia3
+ - AdaLora: https://arxiv.org/abs/2303.10512
+
+ Other PEFT models such as prompt tuning, prompt learning are out of scope as these adapters are not "injectable"
+ into a torch module. For using these methods, please refer to the usage guide of PEFT library.
+
+ With this mixin, if the correct PEFT version is installed, it is possible to:
+
+ - Load an adapter stored on a local path or in a remote Hub repository, and inject it in the model
+ - Attach new adapters in the model and train them with Trainer or by your own.
+ - Attach multiple adapters and iteratively activate / deactivate them
+ - Activate / deactivate all adapters from the model.
+ - Get the `state_dict` of the active adapter.
+ """
+
+ _hf_peft_config_loaded = False
+
+ def load_adapter(
+ self,
+ peft_model_id: Optional[str] = None,
+ adapter_name: Optional[str] = None,
+ revision: Optional[str] = None,
+ token: Optional[str] = None,
+ device_map: Optional[str] = "auto",
+ max_memory: Optional[str] = None,
+ offload_folder: Optional[str] = None,
+ offload_index: Optional[int] = None,
+ peft_config: Dict[str, Any] = None,
+ adapter_state_dict: Optional[Dict[str, "torch.Tensor"]] = None,
+ low_cpu_mem_usage: bool = False,
+ is_trainable: bool = False,
+ adapter_kwargs: Optional[Dict[str, Any]] = None,
+ ) -> None:
+ """
+ Load adapter weights from file or remote Hub folder. If you are not familiar with adapters and PEFT methods, we
+ invite you to read more about them on PEFT official documentation: https://huggingface.co/docs/peft
+
+ Requires peft as a backend to load the adapter weights.
+
+ Args:
+ peft_model_id (`str`, *optional*):
+ The identifier of the model to look for on the Hub, or a local path to the saved adapter config file
+ and adapter weights.
+ adapter_name (`str`, *optional*):
+ The adapter name to use. If not set, will use the default adapter.
+ revision (`str`, *optional*, defaults to `"main"`):
+ The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
+ git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any
+ identifier allowed by git.
+
+
+
+ To test a pull request you made on the Hub, you can pass `revision="refs/pr/"`.
+
+
+
+ token (`str`, `optional`):
+ Whether to use authentication token to load the remote folder. Useful to load private repositories
+ that are on HuggingFace Hub. You might need to call `huggingface-cli login` and paste your tokens to
+ cache it.
+ device_map (`str` or `Dict[str, Union[int, str, torch.device]]` or `int` or `torch.device`, *optional*):
+ A map that specifies where each submodule should go. It doesn't need to be refined to each
+ parameter/buffer name, once a given module name is inside, every submodule of it will be sent to the
+ same device. If we only pass the device (*e.g.*, `"cpu"`, `"cuda:1"`, `"mps"`, or a GPU ordinal rank
+ like `1`) on which the model will be allocated, the device map will map the entire model to this
+ device. Passing `device_map = 0` means put the whole model on GPU 0.
+
+ To have Accelerate compute the most optimized `device_map` automatically, set `device_map="auto"`. For
+ more information about each option see [designing a device
+ map](https://hf.co/docs/accelerate/main/en/usage_guides/big_modeling#designing-a-device-map).
+ max_memory (`Dict`, *optional*):
+ A dictionary device identifier to maximum memory. Will default to the maximum memory available for each
+ GPU and the available CPU RAM if unset.
+ offload_folder (`str` or `os.PathLike`, `optional`):
+ If the `device_map` contains any value `"disk"`, the folder where we will offload weights.
+ offload_index (`int`, `optional`):
+ `offload_index` argument to be passed to `accelerate.dispatch_model` method.
+ peft_config (`Dict[str, Any]`, *optional*):
+ The configuration of the adapter to add, supported adapters are non-prefix tuning and adaption prompts
+ methods. This argument is used in case users directly pass PEFT state dicts
+ adapter_state_dict (`Dict[str, torch.Tensor]`, *optional*):
+ The state dict of the adapter to load. This argument is used in case users directly pass PEFT state
+ dicts
+ low_cpu_mem_usage (`bool`, *optional*, defaults to `False`):
+ Reduce memory usage while loading the PEFT adapter. This should also speed up the loading process.
+ Requires PEFT version 0.13.0 or higher.
+ is_trainable (`bool`, *optional*, defaults to `False`):
+ Whether the adapter should be trainable or not. If `False`, the adapter will be frozen and can only be
+ used for inference.
+ adapter_kwargs (`Dict[str, Any]`, *optional*):
+ Additional keyword arguments passed along to the `from_pretrained` method of the adapter config and
+ `find_adapter_config_file` method.
+ """
+ check_peft_version(min_version=MIN_PEFT_VERSION)
+
+ # peft only supports low_cpu_mem_usage starting from v0.13.0
+ peft_load_kwargs = {}
+ if low_cpu_mem_usage:
+ min_version_lcmu = "0.13.0"
+ if version.parse(importlib.metadata.version("peft")) >= version.parse(min_version_lcmu):
+ peft_load_kwargs["low_cpu_mem_usage"] = low_cpu_mem_usage
+ else:
+ raise ValueError(
+ "The version of PEFT you are using does not support `low_cpu_mem_usage` yet, "
+ f"please install PEFT >= {min_version_lcmu}."
+ )
+
+ adapter_name = adapter_name if adapter_name is not None else "default"
+ if adapter_kwargs is None:
+ adapter_kwargs = {}
+
+ from peft import PeftConfig, inject_adapter_in_model, load_peft_weights
+ from peft.utils import set_peft_model_state_dict
+
+ if self._hf_peft_config_loaded and adapter_name in self.peft_config:
+ raise ValueError(f"Adapter with name {adapter_name} already exists. Please use a different name.")
+
+ if peft_model_id is None and (adapter_state_dict is None and peft_config is None):
+ raise ValueError(
+ "You should either pass a `peft_model_id` or a `peft_config` and `adapter_state_dict` to load an adapter."
+ )
+
+ if "device" not in adapter_kwargs:
+ device = self.device if not hasattr(self, "hf_device_map") else list(self.hf_device_map.values())[0]
+ else:
+ device = adapter_kwargs.pop("device")
+
+ # To avoid PEFT errors later on with safetensors.
+ if isinstance(device, torch.device):
+ device = str(device)
+
+ # We keep `revision` in the signature for backward compatibility
+ if revision is not None and "revision" not in adapter_kwargs:
+ adapter_kwargs["revision"] = revision
+ elif revision is not None and "revision" in adapter_kwargs and revision != adapter_kwargs["revision"]:
+ logger.error(
+ "You passed a `revision` argument both in `adapter_kwargs` and as a standalone argument. "
+ "The one in `adapter_kwargs` will be used."
+ )
+
+ # Override token with adapter_kwargs' token
+ if "token" in adapter_kwargs:
+ token = adapter_kwargs.pop("token")
+
+ if peft_config is None:
+ adapter_config_file = find_adapter_config_file(
+ peft_model_id,
+ token=token,
+ **adapter_kwargs,
+ )
+
+ if adapter_config_file is None:
+ raise ValueError(
+ f"adapter model file not found in {peft_model_id}. Make sure you are passing the correct path to the "
+ "adapter model."
+ )
+
+ peft_config = PeftConfig.from_pretrained(
+ peft_model_id,
+ token=token,
+ **adapter_kwargs,
+ )
+ peft_config.inference_mode = not is_trainable
+
+ # Create and add fresh new adapters into the model.
+ inject_adapter_in_model(peft_config, self, adapter_name, **peft_load_kwargs)
+
+ if not self._hf_peft_config_loaded:
+ self._hf_peft_config_loaded = True
+
+ if peft_model_id is not None:
+ adapter_state_dict = load_peft_weights(peft_model_id, token=token, device=device, **adapter_kwargs)
+
+ # We need to pre-process the state dict to remove unneeded prefixes - for backward compatibility
+ processed_adapter_state_dict = {}
+ prefix = "base_model.model."
+ for key, value in adapter_state_dict.items():
+ if key.startswith(prefix):
+ new_key = key[len(prefix) :]
+ else:
+ new_key = key
+ processed_adapter_state_dict[new_key] = value
+
+ # Load state dict
+ incompatible_keys = set_peft_model_state_dict(
+ self, processed_adapter_state_dict, adapter_name, **peft_load_kwargs
+ )
+
+ if incompatible_keys is not None:
+ err_msg = ""
+ origin_name = peft_model_id if peft_model_id is not None else "state_dict"
+ # Check for unexpected keys.
+ if hasattr(incompatible_keys, "unexpected_keys") and len(incompatible_keys.unexpected_keys) > 0:
+ err_msg = (
+ f"Loading adapter weights from {origin_name} led to unexpected keys not found in the model: "
+ f"{', '.join(incompatible_keys.unexpected_keys)}. "
+ )
+
+ # Check for missing keys.
+ missing_keys = getattr(incompatible_keys, "missing_keys", None)
+ if missing_keys:
+ # Filter missing keys specific to the current adapter, as missing base model keys are expected.
+ lora_missing_keys = [k for k in missing_keys if "lora_" in k and adapter_name in k]
+ if lora_missing_keys:
+ err_msg += (
+ f"Loading adapter weights from {origin_name} led to missing keys in the model: "
+ f"{', '.join(lora_missing_keys)}"
+ )
+
+ if err_msg:
+ logger.warning(err_msg)
+
+ if peft_config.inference_mode:
+ self.eval()
+
+ # Re-dispatch model and hooks in case the model is offloaded to CPU / Disk.
+ if (
+ (getattr(self, "hf_device_map", None) is not None)
+ and (len(set(self.hf_device_map.values()).intersection({"cpu", "disk"})) > 0)
+ and len(self.peft_config) == 1
+ ):
+ self._dispatch_accelerate_model(
+ device_map=device_map,
+ max_memory=max_memory,
+ offload_folder=offload_folder,
+ offload_index=offload_index,
+ )
+
+ def add_adapter(self, adapter_config, adapter_name: Optional[str] = None) -> None:
+ r"""
+ If you are not familiar with adapters and PEFT methods, we invite you to read more about them on the PEFT
+ official documentation: https://huggingface.co/docs/peft
+
+ Adds a fresh new adapter to the current model for training purpose. If no adapter name is passed, a default
+ name is assigned to the adapter to follow the convention of PEFT library (in PEFT we use "default" as the
+ default adapter name).
+
+ Args:
+ adapter_config (`~peft.PeftConfig`):
+ The configuration of the adapter to add, supported adapters are non-prefix tuning and adaption prompts
+ methods
+ adapter_name (`str`, *optional*, defaults to `"default"`):
+ The name of the adapter to add. If no name is passed, a default name is assigned to the adapter.
+ """
+ check_peft_version(min_version=MIN_PEFT_VERSION)
+
+ from peft import PeftConfig, inject_adapter_in_model
+
+ adapter_name = adapter_name or "default"
+
+ if not self._hf_peft_config_loaded:
+ self._hf_peft_config_loaded = True
+ elif adapter_name in self.peft_config:
+ raise ValueError(f"Adapter with name {adapter_name} already exists. Please use a different name.")
+
+ if not isinstance(adapter_config, PeftConfig):
+ raise TypeError(f"adapter_config should be an instance of PeftConfig. Got {type(adapter_config)} instead.")
+
+ # Retrieve the name or path of the model, one could also use self.config._name_or_path
+ # but to be consistent with what we do in PEFT: https://github.com/huggingface/peft/blob/6e783780ca9df3a623992cc4d1d665001232eae0/src/peft/mapping.py#L100
+ adapter_config.base_model_name_or_path = self.__dict__.get("name_or_path", None)
+ inject_adapter_in_model(adapter_config, self, adapter_name)
+
+ self.set_adapter(adapter_name)
+
+ def set_adapter(self, adapter_name: Union[List[str], str]) -> None:
+ """
+ If you are not familiar with adapters and PEFT methods, we invite you to read more about them on the PEFT
+ official documentation: https://huggingface.co/docs/peft
+
+ Sets a specific adapter by forcing the model to use a that adapter and disable the other adapters.
+
+ Args:
+ adapter_name (`Union[List[str], str]`):
+ The name of the adapter to set. Can be also a list of strings to set multiple adapters.
+ """
+ check_peft_version(min_version=MIN_PEFT_VERSION)
+ if not self._hf_peft_config_loaded:
+ raise ValueError("No adapter loaded. Please load an adapter first.")
+ elif isinstance(adapter_name, list):
+ missing = set(adapter_name) - set(self.peft_config)
+ if len(missing) > 0:
+ raise ValueError(
+ f"Following adapter(s) could not be found: {', '.join(missing)}. Make sure you are passing the correct adapter name(s)."
+ f" current loaded adapters are: {list(self.peft_config.keys())}"
+ )
+ elif adapter_name not in self.peft_config:
+ raise ValueError(
+ f"Adapter with name {adapter_name} not found. Please pass the correct adapter name among {list(self.peft_config.keys())}"
+ )
+
+ from peft.tuners.tuners_utils import BaseTunerLayer
+ from peft.utils import ModulesToSaveWrapper
+
+ _adapters_has_been_set = False
+
+ for _, module in self.named_modules():
+ if isinstance(module, (BaseTunerLayer, ModulesToSaveWrapper)):
+ # For backward compatbility with previous PEFT versions
+ if hasattr(module, "set_adapter"):
+ module.set_adapter(adapter_name)
+ else:
+ module.active_adapter = adapter_name
+ _adapters_has_been_set = True
+
+ if not _adapters_has_been_set:
+ raise ValueError(
+ "Did not succeeded in setting the adapter. Please make sure you are using a model that supports adapters."
+ )
+
+ def disable_adapters(self) -> None:
+ r"""
+ If you are not familiar with adapters and PEFT methods, we invite you to read more about them on the PEFT
+ official documentation: https://huggingface.co/docs/peft
+
+ Disable all adapters that are attached to the model. This leads to inferring with the base model only.
+ """
+ check_peft_version(min_version=MIN_PEFT_VERSION)
+
+ if not self._hf_peft_config_loaded:
+ raise ValueError("No adapter loaded. Please load an adapter first.")
+
+ from peft.tuners.tuners_utils import BaseTunerLayer
+ from peft.utils import ModulesToSaveWrapper
+
+ for _, module in self.named_modules():
+ if isinstance(module, (BaseTunerLayer, ModulesToSaveWrapper)):
+ # The recent version of PEFT need to call `enable_adapters` instead
+ if hasattr(module, "enable_adapters"):
+ module.enable_adapters(enabled=False)
+ else:
+ module.disable_adapters = True
+
+ def enable_adapters(self) -> None:
+ """
+ If you are not familiar with adapters and PEFT methods, we invite you to read more about them on the PEFT
+ official documentation: https://huggingface.co/docs/peft
+
+ Enable adapters that are attached to the model.
+ """
+ check_peft_version(min_version=MIN_PEFT_VERSION)
+
+ if not self._hf_peft_config_loaded:
+ raise ValueError("No adapter loaded. Please load an adapter first.")
+
+ from peft.tuners.tuners_utils import BaseTunerLayer
+
+ for _, module in self.named_modules():
+ if isinstance(module, BaseTunerLayer):
+ # The recent version of PEFT need to call `enable_adapters` instead
+ if hasattr(module, "enable_adapters"):
+ module.enable_adapters(enabled=True)
+ else:
+ module.disable_adapters = False
+
+ def active_adapters(self) -> List[str]:
+ """
+ If you are not familiar with adapters and PEFT methods, we invite you to read more about them on the PEFT
+ official documentation: https://huggingface.co/docs/peft
+
+ Gets the current active adapters of the model. In case of multi-adapter inference (combining multiple adapters
+ for inference) returns the list of all active adapters so that users can deal with them accordingly.
+
+ For previous PEFT versions (that does not support multi-adapter inference), `module.active_adapter` will return
+ a single string.
+ """
+ check_peft_version(min_version=MIN_PEFT_VERSION)
+
+ if not is_peft_available():
+ raise ImportError("PEFT is not available. Please install PEFT to use this function: `pip install peft`.")
+
+ if not self._hf_peft_config_loaded:
+ raise ValueError("No adapter loaded. Please load an adapter first.")
+
+ from peft.tuners.tuners_utils import BaseTunerLayer
+
+ for _, module in self.named_modules():
+ if isinstance(module, BaseTunerLayer):
+ active_adapters = module.active_adapter
+ break
+
+ # For previous PEFT versions
+ if isinstance(active_adapters, str):
+ active_adapters = [active_adapters]
+
+ return active_adapters
+
+ def active_adapter(self) -> str:
+ warnings.warn(
+ "The `active_adapter` method is deprecated and will be removed in a future version.", FutureWarning
+ )
+
+ return self.active_adapters()[0]
+
+ def get_adapter_state_dict(self, adapter_name: Optional[str] = None, state_dict: Optional[dict] = None) -> dict:
+ """
+ If you are not familiar with adapters and PEFT methods, we invite you to read more about them on the PEFT
+ official documentation: https://huggingface.co/docs/peft
+
+ Gets the adapter state dict that should only contain the weights tensors of the specified adapter_name adapter.
+ If no adapter_name is passed, the active adapter is used.
+
+ Args:
+ adapter_name (`str`, *optional*):
+ The name of the adapter to get the state dict from. If no name is passed, the active adapter is used.
+ state_dict (nested dictionary of `torch.Tensor`, *optional*)
+ The state dictionary of the model. Will default to `self.state_dict()`, but can be used if special
+ precautions need to be taken when recovering the state dictionary of a model (like when using model
+ parallelism).
+ """
+ check_peft_version(min_version=MIN_PEFT_VERSION)
+
+ if not self._hf_peft_config_loaded:
+ raise ValueError("No adapter loaded. Please load an adapter first.")
+
+ from peft import get_peft_model_state_dict
+
+ if adapter_name is None:
+ adapter_name = self.active_adapters()[0]
+
+ adapter_state_dict = get_peft_model_state_dict(self, state_dict=state_dict, adapter_name=adapter_name)
+ return adapter_state_dict
+
+ def _dispatch_accelerate_model(
+ self,
+ device_map: str,
+ max_memory: Optional[int] = None,
+ offload_folder: Optional[str] = None,
+ offload_index: Optional[int] = None,
+ ) -> None:
+ """
+ Optional re-dispatch the model and attach new hooks to the model in case the model has been loaded with
+ accelerate (i.e. with `device_map=xxx`)
+
+ Args:
+ device_map (`str` or `Dict[str, Union[int, str, torch.device]]` or `int` or `torch.device`, *optional*):
+ A map that specifies where each submodule should go. It doesn't need to be refined to each
+ parameter/buffer name, once a given module name is inside, every submodule of it will be sent to the
+ same device. If we only pass the device (*e.g.*, `"cpu"`, `"cuda:1"`, `"mps"`, or a GPU ordinal rank
+ like `1`) on which the model will be allocated, the device map will map the entire model to this
+ device. Passing `device_map = 0` means put the whole model on GPU 0.
+
+ To have Accelerate compute the most optimized `device_map` automatically, set `device_map="auto"`. For
+ more information about each option see [designing a device
+ map](https://hf.co/docs/accelerate/main/en/usage_guides/big_modeling#designing-a-device-map).
+ max_memory (`Dict`, *optional*):
+ A dictionary device identifier to maximum memory. Will default to the maximum memory available for each
+ GPU and the available CPU RAM if unset.
+ offload_folder (`str` or `os.PathLike`, *optional*):
+ If the `device_map` contains any value `"disk"`, the folder where we will offload weights.
+ offload_index (`int`, *optional*):
+ The offload_index argument to be passed to `accelerate.dispatch_model` method.
+ """
+ dispatch_model_kwargs = {}
+ # Safety checker for previous `accelerate` versions
+ # `offload_index` was introduced in https://github.com/huggingface/accelerate/pull/873/
+ if "offload_index" in inspect.signature(dispatch_model).parameters:
+ dispatch_model_kwargs["offload_index"] = offload_index
+
+ no_split_module_classes = self._no_split_modules
+
+ if device_map != "sequential":
+ max_memory = get_balanced_memory(
+ self,
+ max_memory=max_memory,
+ no_split_module_classes=no_split_module_classes,
+ low_zero=(device_map == "balanced_low_0"),
+ )
+ if isinstance(device_map, str):
+ device_map = infer_auto_device_map(
+ self, max_memory=max_memory, no_split_module_classes=no_split_module_classes
+ )
+ dispatch_model(
+ self,
+ device_map=device_map,
+ offload_dir=offload_folder,
+ **dispatch_model_kwargs,
+ )
+
+ def delete_adapter(self, adapter_names: Union[List[str], str]) -> None:
+ """
+ Delete an adapter's LoRA layers from the underlying model.
+
+ Args:
+ adapter_names (`Union[List[str], str]`):
+ The name(s) of the adapter(s) to delete.
+
+ Example:
+
+ ```py
+ from diffusers import AutoPipelineForText2Image
+ import torch
+
+ pipeline = AutoPipelineForText2Image.from_pretrained(
+ "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16
+ ).to("cuda")
+ pipeline.load_lora_weights(
+ "jbilcke-hf/sdxl-cinematic-1", weight_name="pytorch_lora_weights.safetensors", adapter_names="cinematic"
+ )
+ pipeline.delete_adapters("cinematic")
+ ```
+ """
+
+ check_peft_version(min_version=MIN_PEFT_VERSION)
+
+ if not self._hf_peft_config_loaded:
+ raise ValueError("No adapter loaded. Please load an adapter first.")
+
+ from peft.tuners.tuners_utils import BaseTunerLayer
+
+ if isinstance(adapter_names, str):
+ adapter_names = [adapter_names]
+
+ # Check that all adapter names are present in the config
+ missing_adapters = [name for name in adapter_names if name not in self.peft_config]
+ if missing_adapters:
+ raise ValueError(
+ f"The following adapter(s) are not present and cannot be deleted: {', '.join(missing_adapters)}"
+ )
+
+ for adapter_name in adapter_names:
+ for module in self.modules():
+ if isinstance(module, BaseTunerLayer):
+ if hasattr(module, "delete_adapter"):
+ module.delete_adapter(adapter_name)
+ else:
+ raise ValueError(
+ "The version of PEFT you are using is not compatible, please use a version that is greater than 0.6.1"
+ )
+
+ # For transformers integration - we need to pop the adapter from the config
+ if getattr(self, "_hf_peft_config_loaded", False) and hasattr(self, "peft_config"):
+ self.peft_config.pop(adapter_name, None)
+
+ # In case all adapters are deleted, we need to delete the config
+ # and make sure to set the flag to False
+ if len(self.peft_config) == 0:
+ del self.peft_config
+ self._hf_peft_config_loaded = False
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/peft_utils.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/peft_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..3eb62a099059d4c1acd9eef75f94ee5ac22fd60f
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/peft_utils.py
@@ -0,0 +1,125 @@
+# Copyright 2023 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.
+import importlib
+import os
+from typing import Optional, Union
+
+from packaging import version
+
+from .hub import cached_file
+from .import_utils import is_peft_available
+
+
+ADAPTER_CONFIG_NAME = "adapter_config.json"
+ADAPTER_WEIGHTS_NAME = "adapter_model.bin"
+ADAPTER_SAFE_WEIGHTS_NAME = "adapter_model.safetensors"
+
+
+def find_adapter_config_file(
+ model_id: str,
+ cache_dir: Optional[Union[str, os.PathLike]] = None,
+ force_download: bool = False,
+ resume_download: Optional[bool] = None,
+ proxies: Optional[dict[str, str]] = None,
+ token: Optional[Union[bool, str]] = None,
+ revision: Optional[str] = None,
+ local_files_only: bool = False,
+ subfolder: str = "",
+ _commit_hash: Optional[str] = None,
+) -> Optional[str]:
+ r"""
+ Simply checks if the model stored on the Hub or locally is an adapter model or not, return the path of the adapter
+ config file if it is, None otherwise.
+
+ Args:
+ model_id (`str`):
+ The identifier of the model to look for, can be either a local path or an id to the repository on the Hub.
+ cache_dir (`str` or `os.PathLike`, *optional*):
+ Path to a directory in which a downloaded pretrained model configuration should be cached if the standard
+ cache should not be used.
+ force_download (`bool`, *optional*, defaults to `False`):
+ Whether or not to force to (re-)download the configuration files and override the cached versions if they
+ exist.
+ resume_download:
+ Deprecated and ignored. All downloads are now resumed by default when possible.
+ Will be removed in v5 of Transformers.
+ proxies (`Dict[str, str]`, *optional*):
+ A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',
+ 'http://hostname': 'foo.bar:4012'}.` The proxies are used on each request.
+ token (`str` or *bool*, *optional*):
+ The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated
+ when running `huggingface-cli login` (stored in `~/.huggingface`).
+ revision (`str`, *optional*, defaults to `"main"`):
+ The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
+ git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any
+ identifier allowed by git.
+
+
+
+ To test a pull request you made on the Hub, you can pass `revision="refs/pr/".
+
+
+
+ local_files_only (`bool`, *optional*, defaults to `False`):
+ If `True`, will only try to load the tokenizer configuration from local files.
+ subfolder (`str`, *optional*, defaults to `""`):
+ In case the relevant files are located inside a subfolder of the model repo on huggingface.co, you can
+ specify the folder name here.
+ """
+ adapter_cached_filename = None
+ if model_id is None:
+ return None
+ elif os.path.isdir(model_id):
+ list_remote_files = os.listdir(model_id)
+ if ADAPTER_CONFIG_NAME in list_remote_files:
+ adapter_cached_filename = os.path.join(model_id, ADAPTER_CONFIG_NAME)
+ else:
+ adapter_cached_filename = cached_file(
+ model_id,
+ ADAPTER_CONFIG_NAME,
+ cache_dir=cache_dir,
+ force_download=force_download,
+ resume_download=resume_download,
+ proxies=proxies,
+ token=token,
+ revision=revision,
+ local_files_only=local_files_only,
+ subfolder=subfolder,
+ _commit_hash=_commit_hash,
+ _raise_exceptions_for_gated_repo=False,
+ _raise_exceptions_for_missing_entries=False,
+ _raise_exceptions_for_connection_errors=False,
+ )
+
+ return adapter_cached_filename
+
+
+def check_peft_version(min_version: str) -> None:
+ r"""
+ Checks if the version of PEFT is compatible.
+
+ Args:
+ version (`str`):
+ The version of PEFT to check against.
+ """
+ if not is_peft_available():
+ raise ValueError("PEFT is not installed. Please install it with `pip install peft`")
+
+ is_peft_version_compatible = version.parse(importlib.metadata.version("peft")) >= version.parse(min_version)
+
+ if not is_peft_version_compatible:
+ raise ValueError(
+ f"The version of PEFT you are using is not compatible, please use a version that is greater"
+ f" than {min_version}"
+ )
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/processing_auto.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/processing_auto.py
new file mode 100644
index 0000000000000000000000000000000000000000..6e655edff1c982a9b78ae4dbb60957508d515b46
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/processing_auto.py
@@ -0,0 +1,391 @@
+# coding=utf-8
+# Copyright 2021 The HuggingFace Inc. 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.
+"""AutoProcessor class."""
+
+import importlib
+import inspect
+import json
+import os
+import warnings
+from collections import OrderedDict
+
+# Build the list of all feature extractors
+from ...configuration_utils import PretrainedConfig
+from ...dynamic_module_utils import get_class_from_dynamic_module, resolve_trust_remote_code
+from ...feature_extraction_utils import FeatureExtractionMixin
+from ...image_processing_utils import ImageProcessingMixin
+from ...processing_utils import ProcessorMixin
+from ...tokenization_utils import TOKENIZER_CONFIG_FILE
+from ...utils import FEATURE_EXTRACTOR_NAME, PROCESSOR_NAME, cached_file, logging
+from .auto_factory import _LazyAutoMapping
+from .configuration_auto import (
+ CONFIG_MAPPING_NAMES,
+ AutoConfig,
+ model_type_to_module_name,
+ replace_list_option_in_docstrings,
+)
+from .feature_extraction_auto import AutoFeatureExtractor
+from .image_processing_auto import AutoImageProcessor
+from .tokenization_auto import AutoTokenizer
+
+
+logger = logging.get_logger(__name__)
+
+PROCESSOR_MAPPING_NAMES = OrderedDict(
+ [
+ ("align", "AlignProcessor"),
+ ("altclip", "AltCLIPProcessor"),
+ ("aria", "AriaProcessor"),
+ ("aya_vision", "AyaVisionProcessor"),
+ ("bark", "BarkProcessor"),
+ ("blip", "BlipProcessor"),
+ ("blip-2", "Blip2Processor"),
+ ("bridgetower", "BridgeTowerProcessor"),
+ ("chameleon", "ChameleonProcessor"),
+ ("chinese_clip", "ChineseCLIPProcessor"),
+ ("clap", "ClapProcessor"),
+ ("clip", "CLIPProcessor"),
+ ("clipseg", "CLIPSegProcessor"),
+ ("clvp", "ClvpProcessor"),
+ ("colpali", "ColPaliProcessor"),
+ ("emu3", "Emu3Processor"),
+ ("flava", "FlavaProcessor"),
+ ("fuyu", "FuyuProcessor"),
+ ("gemma3", "Gemma3Processor"),
+ ("git", "GitProcessor"),
+ ("got_ocr2", "GotOcr2Processor"),
+ ("grounding-dino", "GroundingDinoProcessor"),
+ ("groupvit", "CLIPProcessor"),
+ ("hubert", "Wav2Vec2Processor"),
+ ("idefics", "IdeficsProcessor"),
+ ("idefics2", "Idefics2Processor"),
+ ("idefics3", "Idefics3Processor"),
+ ("instructblip", "InstructBlipProcessor"),
+ ("instructblipvideo", "InstructBlipVideoProcessor"),
+ ("kosmos-2", "Kosmos2Processor"),
+ ("layoutlmv2", "LayoutLMv2Processor"),
+ ("layoutlmv3", "LayoutLMv3Processor"),
+ ("llama4", "Llama4Processor"),
+ ("llava", "LlavaProcessor"),
+ ("llava_next", "LlavaNextProcessor"),
+ ("llava_next_video", "LlavaNextVideoProcessor"),
+ ("llava_onevision", "LlavaOnevisionProcessor"),
+ ("markuplm", "MarkupLMProcessor"),
+ ("mctct", "MCTCTProcessor"),
+ ("mgp-str", "MgpstrProcessor"),
+ ("mistral3", "PixtralProcessor"),
+ ("mllama", "MllamaProcessor"),
+ ("moonshine", "Wav2Vec2Processor"),
+ ("oneformer", "OneFormerProcessor"),
+ ("owlv2", "Owlv2Processor"),
+ ("owlvit", "OwlViTProcessor"),
+ ("paligemma", "PaliGemmaProcessor"),
+ ("phi4_multimodal", "Phi4MultimodalProcessor"),
+ ("pix2struct", "Pix2StructProcessor"),
+ ("pixtral", "PixtralProcessor"),
+ ("pop2piano", "Pop2PianoProcessor"),
+ ("qwen2_5_vl", "Qwen2_5_VLProcessor"),
+ ("qwen2_audio", "Qwen2AudioProcessor"),
+ ("qwen2_vl", "Qwen2VLProcessor"),
+ ("sam", "SamProcessor"),
+ ("seamless_m4t", "SeamlessM4TProcessor"),
+ ("sew", "Wav2Vec2Processor"),
+ ("sew-d", "Wav2Vec2Processor"),
+ ("shieldgemma2", "ShieldGemma2Processor"),
+ ("siglip", "SiglipProcessor"),
+ ("siglip2", "Siglip2Processor"),
+ ("speech_to_text", "Speech2TextProcessor"),
+ ("speech_to_text_2", "Speech2Text2Processor"),
+ ("speecht5", "SpeechT5Processor"),
+ ("trocr", "TrOCRProcessor"),
+ ("tvlt", "TvltProcessor"),
+ ("tvp", "TvpProcessor"),
+ ("udop", "UdopProcessor"),
+ ("unispeech", "Wav2Vec2Processor"),
+ ("unispeech-sat", "Wav2Vec2Processor"),
+ ("video_llava", "VideoLlavaProcessor"),
+ ("vilt", "ViltProcessor"),
+ ("vipllava", "LlavaProcessor"),
+ ("vision-text-dual-encoder", "VisionTextDualEncoderProcessor"),
+ ("wav2vec2", "Wav2Vec2Processor"),
+ ("wav2vec2-bert", "Wav2Vec2Processor"),
+ ("wav2vec2-conformer", "Wav2Vec2Processor"),
+ ("wavlm", "Wav2Vec2Processor"),
+ ("whisper", "WhisperProcessor"),
+ ("xclip", "XCLIPProcessor"),
+ ]
+)
+
+PROCESSOR_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, PROCESSOR_MAPPING_NAMES)
+
+
+def processor_class_from_name(class_name: str):
+ for module_name, processors in PROCESSOR_MAPPING_NAMES.items():
+ if class_name in processors:
+ module_name = model_type_to_module_name(module_name)
+
+ module = importlib.import_module(f".{module_name}", "transformers.models")
+ try:
+ return getattr(module, class_name)
+ except AttributeError:
+ continue
+
+ for processor in PROCESSOR_MAPPING._extra_content.values():
+ if getattr(processor, "__name__", None) == class_name:
+ return processor
+
+ # We did not fine the class, but maybe it's because a dep is missing. In that case, the class will be in the main
+ # init and we return the proper dummy to get an appropriate error message.
+ main_module = importlib.import_module("transformers")
+ if hasattr(main_module, class_name):
+ return getattr(main_module, class_name)
+
+ return None
+
+
+class AutoProcessor:
+ r"""
+ This is a generic processor class that will be instantiated as one of the processor classes of the library when
+ created with the [`AutoProcessor.from_pretrained`] class method.
+
+ This class cannot be instantiated directly using `__init__()` (throws an error).
+ """
+
+ def __init__(self):
+ raise EnvironmentError(
+ "AutoProcessor is designed to be instantiated "
+ "using the `AutoProcessor.from_pretrained(pretrained_model_name_or_path)` method."
+ )
+
+ @classmethod
+ @replace_list_option_in_docstrings(PROCESSOR_MAPPING_NAMES)
+ def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+ r"""
+ Instantiate one of the processor classes of the library from a pretrained model vocabulary.
+
+ The processor class to instantiate is selected based on the `model_type` property of the config object (either
+ passed as an argument or loaded from `pretrained_model_name_or_path` if possible):
+
+ List options
+
+ Params:
+ pretrained_model_name_or_path (`str` or `os.PathLike`):
+ This can be either:
+
+ - a string, the *model id* of a pretrained feature_extractor hosted inside a model repo on
+ huggingface.co.
+ - a path to a *directory* containing a processor files saved using the `save_pretrained()` method,
+ e.g., `./my_model_directory/`.
+ cache_dir (`str` or `os.PathLike`, *optional*):
+ Path to a directory in which a downloaded pretrained model feature extractor should be cached if the
+ standard cache should not be used.
+ force_download (`bool`, *optional*, defaults to `False`):
+ Whether or not to force to (re-)download the feature extractor files and override the cached versions
+ if they exist.
+ resume_download:
+ Deprecated and ignored. All downloads are now resumed by default when possible.
+ Will be removed in v5 of Transformers.
+ proxies (`Dict[str, str]`, *optional*):
+ A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',
+ 'http://hostname': 'foo.bar:4012'}.` The proxies are used on each request.
+ token (`str` or *bool*, *optional*):
+ The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated
+ when running `huggingface-cli login` (stored in `~/.huggingface`).
+ revision (`str`, *optional*, defaults to `"main"`):
+ The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
+ git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any
+ identifier allowed by git.
+ return_unused_kwargs (`bool`, *optional*, defaults to `False`):
+ If `False`, then this function returns just the final feature extractor object. If `True`, then this
+ functions returns a `Tuple(feature_extractor, unused_kwargs)` where *unused_kwargs* is a dictionary
+ consisting of the key/value pairs whose keys are not feature extractor attributes: i.e., the part of
+ `kwargs` which has not been used to update `feature_extractor` and is otherwise ignored.
+ trust_remote_code (`bool`, *optional*, defaults to `False`):
+ Whether or not to allow for custom models defined on the Hub in their own modeling files. This option
+ should only be set to `True` for repositories you trust and in which you have read the code, as it will
+ execute code present on the Hub on your local machine.
+ kwargs (`Dict[str, Any]`, *optional*):
+ The values in kwargs of any keys which are feature extractor attributes will be used to override the
+ loaded values. Behavior concerning key/value pairs whose keys are *not* feature extractor attributes is
+ controlled by the `return_unused_kwargs` keyword parameter.
+
+
+
+ Passing `token=True` is required when you want to use a private model.
+
+
+
+ Examples:
+
+ ```python
+ >>> from transformers import AutoProcessor
+
+ >>> # Download processor from huggingface.co and cache.
+ >>> processor = AutoProcessor.from_pretrained("facebook/wav2vec2-base-960h")
+
+ >>> # If processor files are in a directory (e.g. processor was saved using *save_pretrained('./test/saved_model/')*)
+ >>> # processor = AutoProcessor.from_pretrained("./test/saved_model/")
+ ```"""
+ use_auth_token = kwargs.pop("use_auth_token", None)
+ if use_auth_token is not None:
+ warnings.warn(
+ "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.",
+ FutureWarning,
+ )
+ if kwargs.get("token", None) is not None:
+ raise ValueError(
+ "`token` and `use_auth_token` are both specified. Please set only the argument `token`."
+ )
+ kwargs["token"] = use_auth_token
+
+ config = kwargs.pop("config", None)
+ trust_remote_code = kwargs.pop("trust_remote_code", None)
+ kwargs["_from_auto"] = True
+
+ processor_class = None
+ processor_auto_map = None
+
+ # First, let's see if we have a processor or preprocessor config.
+ # Filter the kwargs for `cached_file`.
+ cached_file_kwargs = {
+ key: kwargs[key] for key in inspect.signature(cached_file).parameters.keys() if key in kwargs
+ }
+ # We don't want to raise
+ cached_file_kwargs.update(
+ {
+ "_raise_exceptions_for_gated_repo": False,
+ "_raise_exceptions_for_missing_entries": False,
+ "_raise_exceptions_for_connection_errors": False,
+ }
+ )
+
+ # Let's start by checking whether the processor class is saved in a processor config
+ processor_config_file = cached_file(pretrained_model_name_or_path, PROCESSOR_NAME, **cached_file_kwargs)
+ if processor_config_file is not None:
+ config_dict, _ = ProcessorMixin.get_processor_dict(pretrained_model_name_or_path, **kwargs)
+ processor_class = config_dict.get("processor_class", None)
+ if "AutoProcessor" in config_dict.get("auto_map", {}):
+ processor_auto_map = config_dict["auto_map"]["AutoProcessor"]
+
+ if processor_class is None:
+ # If not found, let's check whether the processor class is saved in an image processor config
+ preprocessor_config_file = cached_file(
+ pretrained_model_name_or_path, FEATURE_EXTRACTOR_NAME, **cached_file_kwargs
+ )
+ if preprocessor_config_file is not None:
+ config_dict, _ = ImageProcessingMixin.get_image_processor_dict(pretrained_model_name_or_path, **kwargs)
+ processor_class = config_dict.get("processor_class", None)
+ if "AutoProcessor" in config_dict.get("auto_map", {}):
+ processor_auto_map = config_dict["auto_map"]["AutoProcessor"]
+
+ # If not found, let's check whether the processor class is saved in a feature extractor config
+ if preprocessor_config_file is not None and processor_class is None:
+ config_dict, _ = FeatureExtractionMixin.get_feature_extractor_dict(
+ pretrained_model_name_or_path, **kwargs
+ )
+ processor_class = config_dict.get("processor_class", None)
+ if "AutoProcessor" in config_dict.get("auto_map", {}):
+ processor_auto_map = config_dict["auto_map"]["AutoProcessor"]
+
+ if processor_class is None:
+ # Next, let's check whether the processor class is saved in a tokenizer
+ tokenizer_config_file = cached_file(
+ pretrained_model_name_or_path, TOKENIZER_CONFIG_FILE, **cached_file_kwargs
+ )
+ if tokenizer_config_file is not None:
+ with open(tokenizer_config_file, encoding="utf-8") as reader:
+ config_dict = json.load(reader)
+
+ processor_class = config_dict.get("processor_class", None)
+ if "AutoProcessor" in config_dict.get("auto_map", {}):
+ processor_auto_map = config_dict["auto_map"]["AutoProcessor"]
+
+ if processor_class is None:
+ # Otherwise, load config, if it can be loaded.
+ if not isinstance(config, PretrainedConfig):
+ config = AutoConfig.from_pretrained(
+ pretrained_model_name_or_path, trust_remote_code=trust_remote_code, **kwargs
+ )
+
+ # And check if the config contains the processor class.
+ processor_class = getattr(config, "processor_class", None)
+ if hasattr(config, "auto_map") and "AutoProcessor" in config.auto_map:
+ processor_auto_map = config.auto_map["AutoProcessor"]
+
+ if processor_class is not None:
+ processor_class = processor_class_from_name(processor_class)
+
+ has_remote_code = processor_auto_map is not None
+ has_local_code = processor_class is not None or type(config) in PROCESSOR_MAPPING
+ trust_remote_code = resolve_trust_remote_code(
+ trust_remote_code, pretrained_model_name_or_path, has_local_code, has_remote_code
+ )
+
+ if has_remote_code and trust_remote_code:
+ processor_class = get_class_from_dynamic_module(
+ processor_auto_map, pretrained_model_name_or_path, **kwargs
+ )
+ _ = kwargs.pop("code_revision", None)
+ if os.path.isdir(pretrained_model_name_or_path):
+ processor_class.register_for_auto_class()
+ return processor_class.from_pretrained(
+ pretrained_model_name_or_path, trust_remote_code=trust_remote_code, **kwargs
+ )
+ elif processor_class is not None:
+ return processor_class.from_pretrained(
+ pretrained_model_name_or_path, trust_remote_code=trust_remote_code, **kwargs
+ )
+ # Last try: we use the PROCESSOR_MAPPING.
+ elif type(config) in PROCESSOR_MAPPING:
+ return PROCESSOR_MAPPING[type(config)].from_pretrained(pretrained_model_name_or_path, **kwargs)
+
+ # At this stage, there doesn't seem to be a `Processor` class available for this model, so let's try a
+ # tokenizer.
+ try:
+ return AutoTokenizer.from_pretrained(
+ pretrained_model_name_or_path, trust_remote_code=trust_remote_code, **kwargs
+ )
+ except Exception:
+ try:
+ return AutoImageProcessor.from_pretrained(
+ pretrained_model_name_or_path, trust_remote_code=trust_remote_code, **kwargs
+ )
+ except Exception:
+ pass
+
+ try:
+ return AutoFeatureExtractor.from_pretrained(
+ pretrained_model_name_or_path, trust_remote_code=trust_remote_code, **kwargs
+ )
+ except Exception:
+ pass
+
+ raise ValueError(
+ f"Unrecognized processing class in {pretrained_model_name_or_path}. Can't instantiate a processor, a "
+ "tokenizer, an image processor or a feature extractor for this model. Make sure the repository contains "
+ "the files of at least one of those processing classes."
+ )
+
+ @staticmethod
+ def register(config_class, processor_class, exist_ok=False):
+ """
+ Register a new processor for this class.
+
+ Args:
+ config_class ([`PretrainedConfig`]):
+ The configuration corresponding to the model to register.
+ processor_class ([`ProcessorMixin`]): The processor to register.
+ """
+ PROCESSOR_MAPPING.register(config_class, processor_class, exist_ok=exist_ok)
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/processing_utils.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/processing_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..b1c40e7ff2d7c08e8b8e741a59f933f58c13fb30
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/processing_utils.py
@@ -0,0 +1,1537 @@
+# Copyright 2022 The HuggingFace Inc. 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.
+"""
+Processing saving/loading class for common processors.
+"""
+
+import copy
+import inspect
+import json
+import os
+import sys
+import typing
+import warnings
+from pathlib import Path
+from typing import Any, Callable, Dict, List, Optional, TypedDict, Union
+
+import numpy as np
+import typing_extensions
+
+from .audio_utils import load_audio
+from .dynamic_module_utils import custom_object_save
+from .image_utils import (
+ ChannelDimension,
+ ImageInput,
+ VideoInput,
+ is_valid_image,
+ is_vision_available,
+ load_image,
+ load_video,
+)
+
+
+if is_vision_available():
+ from .image_utils import PILImageResampling
+
+from .tokenization_utils_base import (
+ PaddingStrategy,
+ PreTokenizedInput,
+ PreTrainedTokenizerBase,
+ TextInput,
+ TruncationStrategy,
+)
+from .utils import (
+ PROCESSOR_NAME,
+ PushToHubMixin,
+ TensorType,
+ add_model_info_to_auto_map,
+ add_model_info_to_custom_pipelines,
+ cached_file,
+ copy_func,
+ direct_transformers_import,
+ download_url,
+ is_offline_mode,
+ is_remote_url,
+ logging,
+)
+
+
+logger = logging.get_logger(__name__)
+
+# Dynamically import the Transformers module to grab the attribute classes of the processor from their names.
+transformers_module = direct_transformers_import(Path(__file__).parent)
+
+
+AUTO_TO_BASE_CLASS_MAPPING = {
+ "AutoTokenizer": "PreTrainedTokenizerBase",
+ "AutoFeatureExtractor": "FeatureExtractionMixin",
+ "AutoImageProcessor": "ImageProcessingMixin",
+}
+
+if sys.version_info >= (3, 11):
+ Unpack = typing.Unpack
+else:
+ Unpack = typing_extensions.Unpack
+
+
+class TextKwargs(TypedDict, total=False):
+ """
+ Keyword arguments for text processing. For extended documentation, check out tokenization_utils_base methods and
+ docstrings associated.
+
+ Attributes:
+ add_special_tokens (`bool`, *optional*)
+ Whether or not to add special tokens when encoding the sequences.
+ padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*)
+ Activates and controls padding.
+ truncation (`bool`, `str` or [`~tokenization_utils_base.TruncationStrategy`], *optional*):
+ Activates and controls truncation.
+ max_length (`int`, *optional*):
+ Controls the maximum length to use by one of the truncation/padding parameters.
+ stride (`int`, *optional*):
+ If set, the overflowing tokens will contain some tokens from the end of the truncated sequence.
+ is_split_into_words (`bool`, *optional*):
+ Whether or not the input is already pre-tokenized.
+ pad_to_multiple_of (`int`, *optional*):
+ If set, will pad the sequence to a multiple of the provided value.
+ return_token_type_ids (`bool`, *optional*):
+ Whether to return token type IDs.
+ return_attention_mask (`bool`, *optional*):
+ Whether to return the attention mask.
+ return_overflowing_tokens (`bool`, *optional*):
+ Whether or not to return overflowing token sequences.
+ return_special_tokens_mask (`bool`, *optional*):
+ Whether or not to return special tokens mask information.
+ return_offsets_mapping (`bool`, *optional*):
+ Whether or not to return `(char_start, char_end)` for each token.
+ return_length (`bool`, *optional*):
+ Whether or not to return the lengths of the encoded inputs.
+ verbose (`bool`, *optional*):
+ Whether or not to print more information and warnings.
+ padding_side (`str`, *optional*):
+ The side on which padding will be applied.
+ """
+
+ text_pair: Optional[Union[TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]]]
+ text_target: Union[TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]]
+ text_pair_target: Optional[Union[TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]]]
+ add_special_tokens: Optional[bool]
+ padding: Union[bool, str, PaddingStrategy]
+ truncation: Union[bool, str, TruncationStrategy]
+ max_length: Optional[int]
+ stride: Optional[int]
+ is_split_into_words: Optional[bool]
+ pad_to_multiple_of: Optional[int]
+ return_token_type_ids: Optional[bool]
+ return_attention_mask: Optional[bool]
+ return_overflowing_tokens: Optional[bool]
+ return_special_tokens_mask: Optional[bool]
+ return_offsets_mapping: Optional[bool]
+ return_length: Optional[bool]
+ verbose: Optional[bool]
+ padding_side: Optional[str]
+
+
+class ImagesKwargs(TypedDict, total=False):
+ """
+ Keyword arguments for image processing. For extended documentation, check the appropriate ImageProcessor
+ class methods and docstrings.
+
+ Attributes:
+ do_resize (`bool`, *optional*):
+ Whether to resize the image.
+ size (`Dict[str, int]`, *optional*):
+ Resize the shorter side of the input to `size["shortest_edge"]`.
+ size_divisor (`int`, *optional*):
+ The size by which to make sure both the height and width can be divided.
+ crop_size (`Dict[str, int]`, *optional*):
+ Desired output size when applying center-cropping.
+ resample (`PILImageResampling`, *optional*):
+ Resampling filter to use if resizing the image.
+ do_rescale (`bool`, *optional*):
+ Whether to rescale the image by the specified scale `rescale_factor`.
+ rescale_factor (`int` or `float`, *optional*):
+ Scale factor to use if rescaling the image.
+ do_normalize (`bool`, *optional*):
+ Whether to normalize the image.
+ image_mean (`float` or `List[float]`, *optional*):
+ Mean to use if normalizing the image.
+ image_std (`float` or `List[float]`, *optional*):
+ Standard deviation to use if normalizing the image.
+ do_pad (`bool`, *optional*):
+ Whether to pad the image to the `(max_height, max_width)` of the images in the batch.
+ pad_size (`Dict[str, int]`, *optional*):
+ The size `{"height": int, "width" int}` to pad the images to.
+ do_center_crop (`bool`, *optional*):
+ Whether to center crop the image.
+ data_format (`ChannelDimension` or `str`, *optional*):
+ The channel dimension format for the output image.
+ input_data_format (`ChannelDimension` or `str`, *optional*):
+ The channel dimension format for the input image.
+ device (`str`, *optional*):
+ The device to use for processing (e.g. "cpu", "cuda"), only relevant for fast image processing.
+ """
+
+ do_resize: Optional[bool]
+ size: Optional[dict[str, int]]
+ size_divisor: Optional[int]
+ crop_size: Optional[dict[str, int]]
+ resample: Optional[Union["PILImageResampling", int]]
+ do_rescale: Optional[bool]
+ rescale_factor: Optional[float]
+ do_normalize: Optional[bool]
+ image_mean: Optional[Union[float, list[float]]]
+ image_std: Optional[Union[float, list[float]]]
+ do_pad: Optional[bool]
+ pad_size: Optional[dict[str, int]]
+ do_center_crop: Optional[bool]
+ data_format: Optional[ChannelDimension]
+ input_data_format: Optional[Union[str, ChannelDimension]]
+ device: Optional[str]
+
+
+class VideosKwargs(TypedDict, total=False):
+ """
+ Keyword arguments for video processing.
+
+ Attributes:
+ do_resize (`bool`):
+ Whether to resize the image.
+ size (`Dict[str, int]`, *optional*):
+ Resize the shorter side of the input to `size["shortest_edge"]`.
+ size_divisor (`int`, *optional*):
+ The size by which to make sure both the height and width can be divided.
+ resample (`PILImageResampling`, *optional*):
+ Resampling filter to use if resizing the image.
+ do_rescale (`bool`, *optional*):
+ Whether to rescale the image by the specified scale `rescale_factor`.
+ rescale_factor (`int` or `float`, *optional*):
+ Scale factor to use if rescaling the image.
+ do_normalize (`bool`, *optional*):
+ Whether to normalize the image.
+ image_mean (`float` or `List[float]`, *optional*):
+ Mean to use if normalizing the image.
+ image_std (`float` or `List[float]`, *optional*):
+ Standard deviation to use if normalizing the image.
+ do_pad (`bool`, *optional*):
+ Whether to pad the image to the `(max_height, max_width)` of the images in the batch.
+ do_center_crop (`bool`, *optional*):
+ Whether to center crop the image.
+ data_format (`ChannelDimension` or `str`, *optional*):
+ The channel dimension format for the output image.
+ input_data_format (`ChannelDimension` or `str`, *optional*):
+ The channel dimension format for the input image.
+ """
+
+ do_resize: Optional[bool]
+ size: Optional[dict[str, int]]
+ size_divisor: Optional[int]
+ resample: Optional["PILImageResampling"]
+ do_rescale: Optional[bool]
+ rescale_factor: Optional[float]
+ do_normalize: Optional[bool]
+ image_mean: Optional[Union[float, list[float]]]
+ image_std: Optional[Union[float, list[float]]]
+ do_pad: Optional[bool]
+ do_center_crop: Optional[bool]
+ data_format: Optional[ChannelDimension]
+ input_data_format: Optional[Union[str, ChannelDimension]]
+
+
+class AudioKwargs(TypedDict, total=False):
+ """
+ Keyword arguments for audio processing.
+
+ Attributes:
+ sampling_rate (`int`, *optional*):
+ The sampling rate at which the `raw_speech` input was sampled.
+ raw_speech (`np.ndarray`, `List[float]`, `List[np.ndarray]`, `List[List[float]]`):
+ The sequence or batch of sequences to be padded. Each sequence can be a numpy array, a list of float
+ values, a list of numpy arrays or a list of list of float values. Must be mono channel audio, not
+ stereo, i.e. single float per timestep.
+ padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*):
+ Select a strategy to pad the returned sequences (according to the model's padding side and padding
+ index) among:
+
+ - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single
+ sequence if provided).
+ - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum
+ acceptable input length for the model if that argument is not provided.
+ - `False` or `'do_not_pad'`
+ max_length (`int`, *optional*):
+ Maximum length of the returned list and optionally padding length (see above).
+ truncation (`bool`, *optional*):
+ Activates truncation to cut input sequences longer than *max_length* to *max_length*.
+ pad_to_multiple_of (`int`, *optional*):
+ If set, will pad the sequence to a multiple of the provided value.
+ return_attention_mask (`bool`, *optional*):
+ Whether or not [`~ASTFeatureExtractor.__call__`] should return `attention_mask`.
+ """
+
+ sampling_rate: Optional[int]
+ raw_speech: Optional[Union["np.ndarray", list[float], list["np.ndarray"], list[list[float]]]]
+ padding: Optional[Union[bool, str, PaddingStrategy]]
+ max_length: Optional[int]
+ truncation: Optional[bool]
+ pad_to_multiple_of: Optional[int]
+ return_attention_mask: Optional[bool]
+
+
+class CommonKwargs(TypedDict, total=False):
+ return_tensors: Optional[Union[str, TensorType]]
+
+
+class ProcessingKwargs(TextKwargs, ImagesKwargs, VideosKwargs, AudioKwargs, CommonKwargs, total=False):
+ """
+ Base class for kwargs passing to processors.
+ A model should have its own `ModelProcessorKwargs` class that inherits from `ProcessingKwargs` to provide:
+ 1) Additional typed keys and that this model requires to process inputs.
+ 2) Default values for existing keys under a `_defaults` attribute.
+ New keys have to be defined as follows to ensure type hinting is done correctly.
+
+ ```python
+ # adding a new image kwarg for this model
+ class ModelImagesKwargs(ImagesKwargs, total=False):
+ new_image_kwarg: Optional[bool]
+
+ class ModelProcessorKwargs(ProcessingKwargs, total=False):
+ images_kwargs: ModelImagesKwargs
+ _defaults = {
+ "images_kwargs: {
+ "new_image_kwarg": False,
+ }
+ "text_kwargs": {
+ "padding": "max_length",
+ },
+ }
+
+ ```
+
+ For Python 3.8 compatibility, when inheriting from this class and overriding one of the kwargs,
+ you need to manually update the __annotations__ dictionary. This can be done as follows:
+
+ ```python
+ class CustomProcessorKwargs(ProcessingKwargs, total=False):
+ images_kwargs: CustomImagesKwargs
+
+ CustomProcessorKwargs.__annotations__["images_kwargs"] = CustomImagesKwargs # python 3.8 compatibility
+ ```python
+
+ """
+
+ common_kwargs: CommonKwargs = {
+ **CommonKwargs.__annotations__,
+ }
+ text_kwargs: TextKwargs = {
+ **TextKwargs.__annotations__,
+ }
+ images_kwargs: ImagesKwargs = {
+ **ImagesKwargs.__annotations__,
+ }
+ videos_kwargs: VideosKwargs = {
+ **VideosKwargs.__annotations__,
+ }
+ audio_kwargs: AudioKwargs = {
+ **AudioKwargs.__annotations__,
+ }
+
+
+class TokenizerChatTemplateKwargs(TypedDict, total=False):
+ """
+ Keyword arguments for tokenizer's `apply_chat_template`, when it is called from within a processor.
+
+ tools (`List[Dict]`, *optional*):
+ A list of tools (callable functions) that will be accessible to the model. If the template does not
+ support function calling, this argument will have no effect. Each tool should be passed as a JSON Schema,
+ giving the name, description and argument types for the tool. See our
+ [chat templating guide](https://huggingface.co/docs/transformers/main/en/chat_templating#automated-function-conversion-for-tool-use)
+ for more information.
+ documents (`List[Dict[str, str]]`, *optional*):
+ A list of dicts representing documents that will be accessible to the model if it is performing RAG
+ (retrieval-augmented generation). If the template does not support RAG, this argument will have no
+ effect. We recommend that each document should be a dict containing "title" and "text" keys. Please
+ see the RAG section of the [chat templating guide](https://huggingface.co/docs/transformers/main/en/chat_templating#arguments-for-RAG)
+ for examples of passing documents with chat templates.
+ add_generation_prompt (bool, *optional*):
+ If this is set, a prompt with the token(s) that indicate
+ the start of an assistant message will be appended to the formatted output. This is useful when you want to generate a response from the model.
+ Note that this argument will be passed to the chat template, and so it must be supported in the
+ template for this argument to have any effect.
+ continue_final_message (bool, *optional*):
+ If this is set, the chat will be formatted so that the final
+ message in the chat is open-ended, without any EOS tokens. The model will continue this message
+ rather than starting a new one. This allows you to "prefill" part of
+ the model's response for it. Cannot be used at the same time as `add_generation_prompt`.
+ return_assistant_tokens_mask (`bool`, defaults to `False`):
+ Whether to return a mask of the assistant generated tokens. For tokens generated by the assistant,
+ the mask will contain 1. For user and system tokens, the mask will contain 0.
+ This functionality is only available for chat templates that support it via the `{% generation %}` keyword.
+ """
+
+ tools: Optional[list[dict]] = None
+ documents: Optional[list[dict[str, str]]] = None
+ add_generation_prompt: Optional[bool] = False
+ continue_final_message: Optional[bool] = False
+ return_assistant_tokens_mask: Optional[bool] = False
+
+
+class ChatTemplateLoadKwargs(TypedDict, total=False):
+ """
+ Keyword arguments used to load multimodal data in processor chat templates.
+
+ num_frames (`int`, *optional*):
+ Number of frames to sample uniformly. If not passed, the whole video is loaded.
+ video_load_backend (`str`, *optional*, defaults to `"pyav"`):
+ The backend to use when loading the video which will be used only when there are videos in the conversation.
+ Can be any of ["decord", "pyav", "opencv", "torchvision"]. Defaults to "pyav" because it is the only backend
+ that supports all types of sources to load from.
+ video_fps (`int`, *optional*):
+ Number of frames to sample per second. Should be passed only when `num_frames=None`.
+ If not specified and `num_frames==None`, all frames are sampled.
+ sample_indices_fn (`Callable`, *optional*):
+ A callable function that will return indices at which the video should be sampled. If the video has to be loaded using
+ by a different sampling technique than provided by `num_frames` or `fps` arguments, one should provide their own `sample_indices_fn`.
+ If not provided, simple uniformt sampling with fps is performed, otherwise `sample_indices_fn` has priority over other args.
+ The function expects at input the all args along with all kwargs passed to `load_video` and should output valid
+ indices at which the video should be sampled. For example:
+
+ def sample_indices_fn(num_frames, fps, metadata, **kwargs):
+ # add you sampling logic here ...
+ return np.linspace(start_idx, end_idx, num_frames, dtype=int)
+ """
+
+ num_frames: Optional[int] = None
+ video_load_backend: Optional[str] = "pyav"
+ video_fps: Optional[int] = None
+ sampling_rate: Optional[int] = 16_000
+ sample_indices_fn: Optional[Callable] = None
+ load_audio_from_video: Optional[bool] = False
+
+
+class ProcessorChatTemplateKwargs(ChatTemplateLoadKwargs, TokenizerChatTemplateKwargs, total=False):
+ """
+ Keyword arguments for processor's `apply_chat_template`.
+
+ tokenize (`bool`, *optional*, defaults to `False`):
+ Whether to tokenize the output or not.
+ return_dict (`bool`, defaults to `False`):
+ Whether to return a dictionary with named outputs. Has no effect if tokenize is `False`.
+ """
+
+ tokenize: Optional[bool] = False
+ return_dict: Optional[bool] = False
+
+
+class AllKwargsForChatTemplate(
+ TextKwargs, ImagesKwargs, VideosKwargs, AudioKwargs, CommonKwargs, ProcessorChatTemplateKwargs
+): ...
+
+
+class ProcessorMixin(PushToHubMixin):
+ """
+ This is a mixin used to provide saving/loading functionality for all processor classes.
+ """
+
+ attributes = ["feature_extractor", "tokenizer"]
+ optional_attributes = ["chat_template"]
+ optional_call_args: list[str] = []
+ # Names need to be attr_class for attr in attributes
+ feature_extractor_class = None
+ tokenizer_class = None
+ _auto_class = None
+ valid_kwargs: list[str] = []
+
+ # args have to match the attributes class attribute
+ def __init__(self, *args, **kwargs):
+ # First, extract optional attributes from kwargs if present
+ # Optional attributes can never be positional arguments
+ for optional_attribute in self.optional_attributes:
+ setattr(self, optional_attribute, kwargs.pop(optional_attribute, None))
+ # Sanitize args and kwargs
+ for key in kwargs:
+ if key not in self.attributes:
+ raise TypeError(f"Unexpected keyword argument {key}.")
+ for arg, attribute_name in zip(args, self.attributes):
+ if attribute_name in kwargs:
+ raise TypeError(f"Got multiple values for argument {attribute_name}.")
+ else:
+ kwargs[attribute_name] = arg
+
+ if len(kwargs) != len(self.attributes):
+ raise ValueError(
+ f"This processor requires {len(self.attributes)} arguments: {', '.join(self.attributes)}. Got "
+ f"{len(args)} arguments instead."
+ )
+
+ # Check each arg is of the proper class (this will also catch a user initializing in the wrong order)
+ for attribute_name, arg in kwargs.items():
+ class_name = getattr(self, f"{attribute_name}_class")
+ # Nothing is ever going to be an instance of "AutoXxx", in that case we check the base class.
+ class_name = AUTO_TO_BASE_CLASS_MAPPING.get(class_name, class_name)
+ if isinstance(class_name, tuple):
+ proper_class = tuple(self.get_possibly_dynamic_module(n) for n in class_name if n is not None)
+ else:
+ proper_class = self.get_possibly_dynamic_module(class_name)
+
+ if not isinstance(arg, proper_class):
+ raise TypeError(
+ f"Received a {type(arg).__name__} for argument {attribute_name}, but a {class_name} was expected."
+ )
+
+ setattr(self, attribute_name, arg)
+
+ def to_dict(self) -> dict[str, Any]:
+ """
+ Serializes this instance to a Python dictionary.
+
+ Returns:
+ `Dict[str, Any]`: Dictionary of all the attributes that make up this processor instance.
+ """
+ output = copy.deepcopy(self.__dict__)
+
+ # Get the kwargs in `__init__`.
+ sig = inspect.signature(self.__init__)
+ # Only save the attributes that are presented in the kwargs of `__init__`.
+ attrs_to_save = sig.parameters
+ # Don't save attributes like `tokenizer`, `image processor` etc.
+ attrs_to_save = [x for x in attrs_to_save if x not in self.__class__.attributes]
+ # extra attributes to be kept
+ attrs_to_save += ["auto_map"]
+
+ output = {k: v for k, v in output.items() if k in attrs_to_save}
+
+ output["processor_class"] = self.__class__.__name__
+
+ if "tokenizer" in output:
+ del output["tokenizer"]
+ if "image_processor" in output:
+ del output["image_processor"]
+ if "feature_extractor" in output:
+ del output["feature_extractor"]
+ if "chat_template" in output:
+ del output["chat_template"]
+
+ # Some attributes have different names but containing objects that are not simple strings
+ output = {
+ k: v
+ for k, v in output.items()
+ if not (isinstance(v, PushToHubMixin) or v.__class__.__name__ == "BeamSearchDecoderCTC")
+ }
+
+ return output
+
+ def to_json_string(self) -> str:
+ """
+ Serializes this instance to a JSON string.
+
+ Returns:
+ `str`: String containing all the attributes that make up this feature_extractor instance in JSON format.
+ """
+ dictionary = self.to_dict()
+
+ return json.dumps(dictionary, indent=2, sort_keys=True) + "\n"
+
+ def to_json_file(self, json_file_path: Union[str, os.PathLike]):
+ """
+ Save this instance to a JSON file.
+
+ Args:
+ json_file_path (`str` or `os.PathLike`):
+ Path to the JSON file in which this processor instance's parameters will be saved.
+ """
+ with open(json_file_path, "w", encoding="utf-8") as writer:
+ writer.write(self.to_json_string())
+
+ def __repr__(self):
+ attributes_repr = [f"- {name}: {repr(getattr(self, name))}" for name in self.attributes]
+ attributes_repr = "\n".join(attributes_repr)
+ return f"{self.__class__.__name__}:\n{attributes_repr}\n\n{self.to_json_string()}"
+
+ def save_pretrained(self, save_directory, push_to_hub: bool = False, **kwargs):
+ """
+ Saves the attributes of this processor (feature extractor, tokenizer...) in the specified directory so that it
+ can be reloaded using the [`~ProcessorMixin.from_pretrained`] method.
+
+
+
+ This class method is simply calling [`~feature_extraction_utils.FeatureExtractionMixin.save_pretrained`] and
+ [`~tokenization_utils_base.PreTrainedTokenizerBase.save_pretrained`]. Please refer to the docstrings of the
+ methods above for more information.
+
+
+
+ Args:
+ save_directory (`str` or `os.PathLike`):
+ Directory where the feature extractor JSON file and the tokenizer files will be saved (directory will
+ be created if it does not exist).
+ push_to_hub (`bool`, *optional*, defaults to `False`):
+ Whether or not to push your model to the Hugging Face model hub after saving it. You can specify the
+ repository you want to push to with `repo_id` (will default to the name of `save_directory` in your
+ namespace).
+ kwargs (`Dict[str, Any]`, *optional*):
+ Additional key word arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method.
+ """
+ use_auth_token = kwargs.pop("use_auth_token", None)
+
+ if use_auth_token is not None:
+ warnings.warn(
+ "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.",
+ FutureWarning,
+ )
+ if kwargs.get("token", None) is not None:
+ raise ValueError(
+ "`token` and `use_auth_token` are both specified. Please set only the argument `token`."
+ )
+ kwargs["token"] = use_auth_token
+
+ os.makedirs(save_directory, exist_ok=True)
+
+ if push_to_hub:
+ commit_message = kwargs.pop("commit_message", None)
+ repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1])
+ repo_id = self._create_repo(repo_id, **kwargs)
+ files_timestamps = self._get_files_timestamps(save_directory)
+ # If we have a custom config, we copy the file defining it in the folder and set the attributes so it can be
+ # loaded from the Hub.
+ if self._auto_class is not None:
+ attrs = [getattr(self, attribute_name) for attribute_name in self.attributes]
+ configs = [(a.init_kwargs if isinstance(a, PreTrainedTokenizerBase) else a) for a in attrs]
+ configs.append(self)
+ custom_object_save(self, save_directory, config=configs)
+
+ for attribute_name in self.attributes:
+ attribute = getattr(self, attribute_name)
+ # Include the processor class in the attribute config so this processor can then be reloaded with the
+ # `AutoProcessor` API.
+ if hasattr(attribute, "_set_processor_class"):
+ attribute._set_processor_class(self.__class__.__name__)
+ attribute.save_pretrained(save_directory)
+
+ if self._auto_class is not None:
+ # We added an attribute to the init_kwargs of the tokenizers, which needs to be cleaned up.
+ for attribute_name in self.attributes:
+ attribute = getattr(self, attribute_name)
+ if isinstance(attribute, PreTrainedTokenizerBase):
+ del attribute.init_kwargs["auto_map"]
+
+ # If we save using the predefined names, we can load using `from_pretrained`
+ # plus we save chat_template in its own file
+ output_processor_file = os.path.join(save_directory, PROCESSOR_NAME)
+ output_raw_chat_template_file = os.path.join(save_directory, "chat_template.jinja")
+ output_chat_template_file = os.path.join(save_directory, "chat_template.json")
+
+ processor_dict = self.to_dict()
+ # Save `chat_template` in its own file. We can't get it from `processor_dict` as we popped it in `to_dict`
+ # to avoid serializing chat template in json config file. So let's get it from `self` directly
+ if self.chat_template is not None:
+ if kwargs.get("save_raw_chat_template", False):
+ with open(output_raw_chat_template_file, "w", encoding="utf-8") as writer:
+ writer.write(self.chat_template)
+ logger.info(f"chat template saved in {output_raw_chat_template_file}")
+ else:
+ chat_template_json_string = (
+ json.dumps({"chat_template": self.chat_template}, indent=2, sort_keys=True) + "\n"
+ )
+ with open(output_chat_template_file, "w", encoding="utf-8") as writer:
+ writer.write(chat_template_json_string)
+ logger.info(f"chat template saved in {output_chat_template_file}")
+
+ # For now, let's not save to `processor_config.json` if the processor doesn't have extra attributes and
+ # `auto_map` is not specified.
+ if set(processor_dict.keys()) != {"processor_class"}:
+ self.to_json_file(output_processor_file)
+ logger.info(f"processor saved in {output_processor_file}")
+
+ if push_to_hub:
+ self._upload_modified_files(
+ save_directory,
+ repo_id,
+ files_timestamps,
+ commit_message=commit_message,
+ token=kwargs.get("token"),
+ )
+
+ if set(processor_dict.keys()) == {"processor_class"}:
+ return []
+ return [output_processor_file]
+
+ @classmethod
+ def get_processor_dict(
+ cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs
+ ) -> tuple[dict[str, Any], dict[str, Any]]:
+ """
+ From a `pretrained_model_name_or_path`, resolve to a dictionary of parameters, to be used for instantiating a
+ processor of type [`~processing_utils.ProcessingMixin`] using `from_args_and_dict`.
+
+ Parameters:
+ pretrained_model_name_or_path (`str` or `os.PathLike`):
+ The identifier of the pre-trained checkpoint from which we want the dictionary of parameters.
+ subfolder (`str`, *optional*, defaults to `""`):
+ In case the relevant files are located inside a subfolder of the model repo on huggingface.co, you can
+ specify the folder name here.
+
+ Returns:
+ `Tuple[Dict, Dict]`: The dictionary(ies) that will be used to instantiate the processor object.
+ """
+ cache_dir = kwargs.pop("cache_dir", None)
+ force_download = kwargs.pop("force_download", False)
+ resume_download = kwargs.pop("resume_download", None)
+ proxies = kwargs.pop("proxies", None)
+ token = kwargs.pop("token", None)
+ local_files_only = kwargs.pop("local_files_only", False)
+ revision = kwargs.pop("revision", None)
+ subfolder = kwargs.pop("subfolder", "")
+
+ from_pipeline = kwargs.pop("_from_pipeline", None)
+ from_auto_class = kwargs.pop("_from_auto", False)
+
+ user_agent = {"file_type": "processor", "from_auto_class": from_auto_class}
+ if from_pipeline is not None:
+ user_agent["using_pipeline"] = from_pipeline
+
+ if is_offline_mode() and not local_files_only:
+ logger.info("Offline mode: forcing local_files_only=True")
+ local_files_only = True
+
+ pretrained_model_name_or_path = str(pretrained_model_name_or_path)
+ is_local = os.path.isdir(pretrained_model_name_or_path)
+ if os.path.isdir(pretrained_model_name_or_path):
+ processor_file = os.path.join(pretrained_model_name_or_path, PROCESSOR_NAME)
+
+ if os.path.isfile(pretrained_model_name_or_path):
+ resolved_processor_file = pretrained_model_name_or_path
+ # cant't load chat-template when given a file as pretrained_model_name_or_path
+ resolved_chat_template_file = None
+ resolved_raw_chat_template_file = None
+ is_local = True
+ elif is_remote_url(pretrained_model_name_or_path):
+ processor_file = pretrained_model_name_or_path
+ resolved_processor_file = download_url(pretrained_model_name_or_path)
+ # can't load chat-template when given a file url as pretrained_model_name_or_path
+ resolved_chat_template_file = None
+ resolved_raw_chat_template_file = None
+ else:
+ processor_file = PROCESSOR_NAME
+ chat_template_file = "chat_template.json"
+ raw_chat_template_file = "chat_template.jinja"
+ try:
+ # Load from local folder or from cache or download from model Hub and cache
+ resolved_processor_file = cached_file(
+ pretrained_model_name_or_path,
+ processor_file,
+ cache_dir=cache_dir,
+ force_download=force_download,
+ proxies=proxies,
+ resume_download=resume_download,
+ local_files_only=local_files_only,
+ token=token,
+ user_agent=user_agent,
+ revision=revision,
+ subfolder=subfolder,
+ _raise_exceptions_for_missing_entries=False,
+ )
+
+ # Load chat template from a separate json if exists
+ # because making it part of processor-config break BC.
+ # Processors in older version do not accept any kwargs
+ resolved_chat_template_file = cached_file(
+ pretrained_model_name_or_path,
+ chat_template_file,
+ cache_dir=cache_dir,
+ force_download=force_download,
+ proxies=proxies,
+ resume_download=resume_download,
+ local_files_only=local_files_only,
+ token=token,
+ user_agent=user_agent,
+ revision=revision,
+ subfolder=subfolder,
+ _raise_exceptions_for_missing_entries=False,
+ )
+
+ resolved_raw_chat_template_file = cached_file(
+ pretrained_model_name_or_path,
+ raw_chat_template_file,
+ cache_dir=cache_dir,
+ force_download=force_download,
+ proxies=proxies,
+ resume_download=resume_download,
+ local_files_only=local_files_only,
+ token=token,
+ user_agent=user_agent,
+ revision=revision,
+ subfolder=subfolder,
+ _raise_exceptions_for_missing_entries=False,
+ )
+ except OSError:
+ # Raise any environment error raise by `cached_file`. It will have a helpful error message adapted to
+ # the original exception.
+ raise
+ except Exception:
+ # For any other exception, we throw a generic error.
+ raise OSError(
+ f"Can't load processor for '{pretrained_model_name_or_path}'. If you were trying to load"
+ " it from 'https://huggingface.co/models', make sure you don't have a local directory with the"
+ f" same name. Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a"
+ f" directory containing a {PROCESSOR_NAME} file"
+ )
+
+ # Add chat template as kwarg before returning because most models don't have processor config
+ if resolved_raw_chat_template_file is not None:
+ with open(resolved_raw_chat_template_file, encoding="utf-8") as reader:
+ chat_template = reader.read()
+ kwargs["chat_template"] = chat_template
+ elif resolved_chat_template_file is not None:
+ with open(resolved_chat_template_file, encoding="utf-8") as reader:
+ text = reader.read()
+ chat_template = json.loads(text)["chat_template"]
+ kwargs["chat_template"] = chat_template
+
+ # Existing processors on the Hub created before #27761 being merged don't have `processor_config.json` (if not
+ # updated afterward), and we need to keep `from_pretrained` work. So here it fallbacks to the empty dict.
+ # (`cached_file` called using `_raise_exceptions_for_missing_entries=False` to avoid exception)
+ # However, for models added in the future, we won't get the expected error if this file is missing.
+ if resolved_processor_file is None:
+ # In any case we need to pass `chat_template` if it is available
+ processor_dict = {}
+ if "chat_template" in kwargs:
+ processor_dict = {"chat_template": kwargs.pop("chat_template")}
+ return processor_dict, kwargs
+
+ try:
+ # Load processor dict
+ with open(resolved_processor_file, encoding="utf-8") as reader:
+ text = reader.read()
+ processor_dict = json.loads(text)
+
+ except json.JSONDecodeError:
+ raise OSError(f"It looks like the config file at '{resolved_processor_file}' is not a valid JSON file.")
+
+ if is_local:
+ logger.info(f"loading configuration file {resolved_processor_file}")
+ else:
+ logger.info(f"loading configuration file {processor_file} from cache at {resolved_processor_file}")
+
+ if "chat_template" in processor_dict and processor_dict["chat_template"] is not None:
+ logger.warning_once(
+ "Chat templates should be in a 'chat_template.jinja' file but found key='chat_template' "
+ "in the processor's config. Make sure to move your template to its own file."
+ )
+
+ if "chat_template" in kwargs:
+ processor_dict["chat_template"] = kwargs.pop("chat_template")
+
+ if not is_local:
+ if "auto_map" in processor_dict:
+ processor_dict["auto_map"] = add_model_info_to_auto_map(
+ processor_dict["auto_map"], pretrained_model_name_or_path
+ )
+ if "custom_pipelines" in processor_dict:
+ processor_dict["custom_pipelines"] = add_model_info_to_custom_pipelines(
+ processor_dict["custom_pipelines"], pretrained_model_name_or_path
+ )
+
+ return processor_dict, kwargs
+
+ @classmethod
+ def from_args_and_dict(cls, args, processor_dict: dict[str, Any], **kwargs):
+ """
+ Instantiates a type of [`~processing_utils.ProcessingMixin`] from a Python dictionary of parameters.
+
+ Args:
+ processor_dict (`Dict[str, Any]`):
+ Dictionary that will be used to instantiate the processor object. Such a dictionary can be
+ retrieved from a pretrained checkpoint by leveraging the
+ [`~processing_utils.ProcessingMixin.to_dict`] method.
+ kwargs (`Dict[str, Any]`):
+ Additional parameters from which to initialize the processor object.
+
+ Returns:
+ [`~processing_utils.ProcessingMixin`]: The processor object instantiated from those
+ parameters.
+ """
+ processor_dict = processor_dict.copy()
+ return_unused_kwargs = kwargs.pop("return_unused_kwargs", False)
+
+ # We have to pop up some unused (but specific) kwargs and then validate that it doesn't contain unused kwargs
+ # If we don't pop, some specific kwargs will raise a warning
+ if "processor_class" in processor_dict:
+ del processor_dict["processor_class"]
+
+ if "auto_map" in processor_dict:
+ del processor_dict["auto_map"]
+
+ unused_kwargs = cls.validate_init_kwargs(processor_config=processor_dict, valid_kwargs=cls.valid_kwargs)
+ processor = cls(*args, **processor_dict)
+
+ # Update processor with kwargs if needed
+ for key in set(kwargs.keys()):
+ if hasattr(processor, key):
+ setattr(processor, key, kwargs.pop(key))
+
+ kwargs.update(unused_kwargs)
+ logger.info(f"Processor {processor}")
+ if return_unused_kwargs:
+ return processor, kwargs
+ else:
+ return processor
+
+ def _merge_kwargs(
+ self,
+ ModelProcessorKwargs: ProcessingKwargs,
+ tokenizer_init_kwargs: Optional[dict] = None,
+ **kwargs,
+ ) -> dict[str, dict]:
+ """
+ Method to merge dictionaries of kwargs cleanly separated by modality within a Processor instance.
+ The order of operations is as follows:
+ 1) kwargs passed as before have highest priority to preserve BC.
+ ```python
+ high_priority_kwargs = {"crop_size" = {"height": 222, "width": 222}, "padding" = "max_length"}
+ processor(..., **high_priority_kwargs)
+ ```
+ 2) kwargs passed as modality-specific kwargs have second priority. This is the recommended API.
+ ```python
+ processor(..., text_kwargs={"padding": "max_length"}, images_kwargs={"crop_size": {"height": 222, "width": 222}}})
+ ```
+ 3) kwargs passed during instantiation of a modality processor have fourth priority.
+ ```python
+ tokenizer = tokenizer_class(..., {"padding": "max_length"})
+ image_processor = image_processor_class(...)
+ processor(tokenizer, image_processor) # will pass max_length unless overridden by kwargs at call
+ ```
+ 4) defaults kwargs specified at processor level have lowest priority.
+ ```python
+ class MyProcessingKwargs(ProcessingKwargs, CommonKwargs, TextKwargs, ImagesKwargs, total=False):
+ _defaults = {
+ "text_kwargs": {
+ "padding": "max_length",
+ "max_length": 64,
+ },
+ }
+ ```
+ Args:
+ ModelProcessorKwargs (`ProcessingKwargs`):
+ Typed dictionary of kwargs specifically required by the model passed.
+ tokenizer_init_kwargs (`Dict`, *optional*):
+ Dictionary of kwargs the tokenizer was instantiated with and need to take precedence over defaults.
+
+ Returns:
+ output_kwargs (`Dict`):
+ Dictionary of per-modality kwargs to be passed to each modality-specific processor.
+
+ """
+ # Initialize dictionaries
+ output_kwargs = {
+ "text_kwargs": {},
+ "images_kwargs": {},
+ "audio_kwargs": {},
+ "videos_kwargs": {},
+ "common_kwargs": {},
+ }
+
+ default_kwargs = {
+ "text_kwargs": {},
+ "images_kwargs": {},
+ "audio_kwargs": {},
+ "videos_kwargs": {},
+ "common_kwargs": {},
+ }
+
+ possible_modality_keywords = {"text", "audio", "videos", "images"}
+ used_keys = set()
+
+ # get defaults from set model processor kwargs if they exist
+ for modality in default_kwargs:
+ default_kwargs[modality] = ModelProcessorKwargs._defaults.get(modality, {}).copy()
+ # update defaults with arguments from tokenizer init
+ for modality_key in ModelProcessorKwargs.__annotations__[modality].__annotations__.keys():
+ # init with tokenizer init kwargs if necessary
+ if modality_key in tokenizer_init_kwargs:
+ value = (
+ getattr(self.tokenizer, modality_key)
+ if hasattr(self.tokenizer, modality_key)
+ else tokenizer_init_kwargs[modality_key]
+ )
+ default_kwargs[modality][modality_key] = value
+ # now defaults kwargs are updated with the tokenizers defaults.
+ # pass defaults to output dictionary
+ output_kwargs.update(default_kwargs)
+
+ # update modality kwargs with passed kwargs
+ non_modality_kwargs = set(kwargs) - set(output_kwargs)
+ for modality in output_kwargs:
+ for modality_key in ModelProcessorKwargs.__annotations__[modality].__annotations__.keys():
+ # check if we received a structured kwarg dict or not to handle it correctly
+ if modality in kwargs:
+ kwarg_value = kwargs[modality].pop(modality_key, "__empty__")
+ # check if this key was passed as a flat kwarg.
+ if kwarg_value != "__empty__" and modality_key in non_modality_kwargs:
+ raise ValueError(
+ f"Keyword argument {modality_key} was passed two times:\n"
+ f"in a dictionary for {modality} and as a **kwarg."
+ )
+ elif modality_key in kwargs:
+ # we get a modality_key instead of popping it because modality-specific processors
+ # can have overlapping kwargs
+ kwarg_value = kwargs.get(modality_key, "__empty__")
+ else:
+ kwarg_value = "__empty__"
+ if not isinstance(kwarg_value, str) or kwarg_value != "__empty__":
+ output_kwargs[modality][modality_key] = kwarg_value
+ used_keys.add(modality_key)
+
+ # Determine if kwargs is a flat dictionary or contains nested dictionaries
+ if any(key in default_kwargs for key in kwargs):
+ # kwargs is dictionary-based, and some keys match modality names
+ for modality, subdict in kwargs.items():
+ if modality in default_kwargs:
+ for subkey, subvalue in subdict.items():
+ if subkey not in used_keys:
+ output_kwargs[modality][subkey] = subvalue
+ used_keys.add(subkey)
+ else:
+ # kwargs is a flat dictionary
+ for key in kwargs:
+ if key not in used_keys:
+ if key in ModelProcessorKwargs.__annotations__["common_kwargs"].__annotations__.keys():
+ output_kwargs["common_kwargs"][key] = kwargs[key]
+ elif key not in possible_modality_keywords:
+ logger.warning_once(
+ f"Keyword argument `{key}` is not a valid argument for this processor and will be ignored."
+ )
+
+ # all modality-specific kwargs are updated with common kwargs
+ for modality in output_kwargs:
+ output_kwargs[modality].update(output_kwargs["common_kwargs"])
+ return output_kwargs
+
+ @classmethod
+ def from_pretrained(
+ cls,
+ pretrained_model_name_or_path: Union[str, os.PathLike],
+ cache_dir: Optional[Union[str, os.PathLike]] = None,
+ force_download: bool = False,
+ local_files_only: bool = False,
+ token: Optional[Union[str, bool]] = None,
+ revision: str = "main",
+ **kwargs,
+ ):
+ r"""
+ Instantiate a processor associated with a pretrained model.
+
+
+
+ This class method is simply calling the feature extractor
+ [`~feature_extraction_utils.FeatureExtractionMixin.from_pretrained`], image processor
+ [`~image_processing_utils.ImageProcessingMixin`] and the tokenizer
+ [`~tokenization_utils_base.PreTrainedTokenizer.from_pretrained`] methods. Please refer to the docstrings of the
+ methods above for more information.
+
+
+
+ Args:
+ pretrained_model_name_or_path (`str` or `os.PathLike`):
+ This can be either:
+
+ - a string, the *model id* of a pretrained feature_extractor hosted inside a model repo on
+ huggingface.co.
+ - a path to a *directory* containing a feature extractor file saved using the
+ [`~SequenceFeatureExtractor.save_pretrained`] method, e.g., `./my_model_directory/`.
+ - a path or url to a saved feature extractor JSON *file*, e.g.,
+ `./my_model_directory/preprocessor_config.json`.
+ **kwargs
+ Additional keyword arguments passed along to both
+ [`~feature_extraction_utils.FeatureExtractionMixin.from_pretrained`] and
+ [`~tokenization_utils_base.PreTrainedTokenizer.from_pretrained`].
+ """
+ kwargs["cache_dir"] = cache_dir
+ kwargs["force_download"] = force_download
+ kwargs["local_files_only"] = local_files_only
+ kwargs["revision"] = revision
+
+ use_auth_token = kwargs.pop("use_auth_token", None)
+ if use_auth_token is not None:
+ warnings.warn(
+ "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.",
+ FutureWarning,
+ )
+ if token is not None:
+ raise ValueError(
+ "`token` and `use_auth_token` are both specified. Please set only the argument `token`."
+ )
+ token = use_auth_token
+
+ if token is not None:
+ kwargs["token"] = token
+
+ args = cls._get_arguments_from_pretrained(pretrained_model_name_or_path, **kwargs)
+ processor_dict, kwargs = cls.get_processor_dict(pretrained_model_name_or_path, **kwargs)
+ processor_dict.update({k: v for k, v in kwargs.items() if k in processor_dict.keys()})
+ return cls.from_args_and_dict(args, processor_dict, **kwargs)
+
+ @classmethod
+ def register_for_auto_class(cls, auto_class="AutoProcessor"):
+ """
+ Register this class with a given auto class. This should only be used for custom feature extractors as the ones
+ in the library are already mapped with `AutoProcessor`.
+
+
+
+ This API is experimental and may have some slight breaking changes in the next releases.
+
+
+
+ Args:
+ auto_class (`str` or `type`, *optional*, defaults to `"AutoProcessor"`):
+ The auto class to register this new feature extractor with.
+ """
+ if not isinstance(auto_class, str):
+ auto_class = auto_class.__name__
+
+ import transformers.models.auto as auto_module
+
+ if not hasattr(auto_module, auto_class):
+ raise ValueError(f"{auto_class} is not a valid auto class.")
+
+ cls._auto_class = auto_class
+
+ @classmethod
+ def _get_arguments_from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+ """
+ Identify and instantiate the subcomponents of Processor classes, like image processors and
+ tokenizers. This method uses the Processor attributes like `tokenizer_class` to figure out what class those
+ subcomponents should be. Note that any subcomponents must either be library classes that are accessible in
+ the `transformers` root, or they must be custom code that has been registered with the relevant autoclass,
+ via methods like `AutoTokenizer.register()`. If neither of these conditions are fulfilled, this method
+ will be unable to find the relevant subcomponent class and will raise an error.
+ """
+ args = []
+ for attribute_name in cls.attributes:
+ class_name = getattr(cls, f"{attribute_name}_class")
+ if isinstance(class_name, tuple):
+ classes = tuple(cls.get_possibly_dynamic_module(n) if n is not None else None for n in class_name)
+ if attribute_name == "image_processor":
+ # TODO: @yoni, change logic in v4.52 (when use_fast set to True by default)
+ use_fast = kwargs.get("use_fast", None)
+ if use_fast is None:
+ logger.warning_once(
+ "Using a slow image processor as `use_fast` is unset and a slow processor was saved with this model. "
+ "`use_fast=True` will be the default behavior in v4.52, even if the model was saved with a slow processor. "
+ "This will result in minor differences in outputs. You'll still be able to use a slow processor with `use_fast=False`."
+ )
+ else:
+ use_fast = kwargs.get("use_fast", True)
+ if use_fast and classes[1] is not None:
+ attribute_class = classes[1]
+ else:
+ attribute_class = classes[0]
+ else:
+ attribute_class = cls.get_possibly_dynamic_module(class_name)
+
+ args.append(attribute_class.from_pretrained(pretrained_model_name_or_path, **kwargs))
+ return args
+
+ @staticmethod
+ def get_possibly_dynamic_module(module_name):
+ if hasattr(transformers_module, module_name):
+ return getattr(transformers_module, module_name)
+ lookup_locations = [
+ transformers_module.IMAGE_PROCESSOR_MAPPING,
+ transformers_module.TOKENIZER_MAPPING,
+ transformers_module.FEATURE_EXTRACTOR_MAPPING,
+ ]
+ for lookup_location in lookup_locations:
+ for custom_class in lookup_location._extra_content.values():
+ if isinstance(custom_class, tuple):
+ for custom_subclass in custom_class:
+ if custom_subclass is not None and custom_subclass.__name__ == module_name:
+ return custom_subclass
+ elif custom_class is not None and custom_class.__name__ == module_name:
+ return custom_class
+ else:
+ raise ValueError(
+ f"Could not find module {module_name} in `transformers`. If this is a custom class, "
+ f"it should be registered using the relevant `AutoClass.register()` function so that "
+ f"other functions can find it!"
+ )
+
+ @property
+ def model_input_names(self):
+ first_attribute = getattr(self, self.attributes[0])
+ return getattr(first_attribute, "model_input_names", None)
+
+ @staticmethod
+ def validate_init_kwargs(processor_config, valid_kwargs):
+ kwargs_from_config = processor_config.keys()
+ unused_kwargs = {}
+ unused_keys = set(kwargs_from_config) - set(valid_kwargs)
+ if unused_keys:
+ unused_kwargs = {k: processor_config[k] for k in unused_keys}
+ return unused_kwargs
+
+ def prepare_and_validate_optional_call_args(self, *args):
+ """
+ Matches optional positional arguments to their corresponding names in `optional_call_args`
+ in the processor class in the order they are passed to the processor call.
+
+ Note that this should only be used in the `__call__` method of the processors with special
+ arguments. Special arguments are arguments that aren't `text`, `images`, `audio`, nor `videos`
+ but also aren't passed to the tokenizer, image processor, etc. Examples of such processors are:
+ - `CLIPSegProcessor`
+ - `LayoutLMv2Processor`
+ - `OwlViTProcessor`
+
+ Also note that passing by position to the processor call is now deprecated and will be disallowed
+ in future versions. We only have this for backward compatibility.
+
+ Example:
+ Suppose that the processor class has `optional_call_args = ["arg_name_1", "arg_name_2"]`.
+ And we define the call method as:
+ ```python
+ def __call__(
+ self,
+ text: str,
+ images: Optional[ImageInput] = None,
+ *arg,
+ audio=None,
+ videos=None,
+ )
+ ```
+
+ Then, if we call the processor as:
+ ```python
+ images = [...]
+ processor("What is common in these images?", images, arg_value_1, arg_value_2)
+ ```
+
+ Then, this method will return:
+ ```python
+ {
+ "arg_name_1": arg_value_1,
+ "arg_name_2": arg_value_2,
+ }
+ ```
+ which we could then pass as kwargs to `self._merge_kwargs`
+ """
+ if len(args):
+ warnings.warn(
+ "Passing positional arguments to the processor call is now deprecated and will be disallowed in v4.47. "
+ "Please pass all arguments as keyword arguments."
+ )
+ if len(args) > len(self.optional_call_args):
+ raise ValueError(
+ f"Expected *at most* {len(self.optional_call_args)} optional positional arguments in processor call"
+ f"which will be matched with {' '.join(self.optional_call_args)} in the order they are passed."
+ f"However, got {len(args)} positional arguments instead."
+ "Please pass all arguments as keyword arguments instead (e.g. `processor(arg_name_1=..., arg_name_2=...))`."
+ )
+ return {arg_name: arg_value for arg_value, arg_name in zip(args, self.optional_call_args)}
+
+ def _process_messages_for_chat_template(
+ self,
+ conversation: List[List[Dict[str, str]]],
+ batch_images: List[ImageInput],
+ batch_videos: List[VideoInput],
+ batch_video_metadata: List[List[Dict[str, any]]],
+ **mm_load_kwargs: Unpack[ChatTemplateLoadKwargs],
+ ):
+ """
+ Used within `apply_chat_template` when a model has a special way to process conversation history. For example,
+ video models might want to specify in the prompt the duration of video or which frame indices at which timestamps
+ were sampled. This information cannot be accessed before the video is loaded.
+
+ For most models it is a no-op, and must be overridden by model processors which require special processing.
+
+ Args:
+ conversation (`List[Dict, str, str]`):
+ The conversation to process. Always comes in batched format.
+ batch_images (`List[List[ImageInput]]`):
+ Batch of images that were loaded from url/path defined in the conversation. The images
+ are ordered in the same way as in the conversation. Comes in nested list format, one list of `PIL` images
+ per batch.
+ batch_videos (`List[List[ImageInput]]`):
+ Batch of videos that were loaded from url/path defined in the conversation. The videos
+ are ordered in the samm way as in the conversation. Comes in nested list format, one list of 4D video arrays
+ per batch.
+ batch_video_metadata (`List[List[Dict[[str, any]]]]`):
+ Batch of metadata returned from loading videos. That includes video fps, duration and total number of framer in original video.
+ Metadata are ordered in the same way as `batch_videos`. Comes in nested list format, one list of 4D video arrays
+ per batch.
+
+ """
+ return conversation
+
+ def apply_chat_template(
+ self,
+ conversation: Union[list[dict[str, str]], list[list[dict[str, str]]]],
+ chat_template: Optional[str] = None,
+ **kwargs: Unpack[AllKwargsForChatTemplate],
+ ) -> str:
+ """
+ Similar to the `apply_chat_template` method on tokenizers, this method applies a Jinja template to input
+ conversations to turn them into a single tokenizable string.
+
+ The input is expected to be in the following format, where each message content is a list consisting of text and
+ optionally image or video inputs. One can also provide an image, video, URL or local path which will be used to form
+ `pixel_values` when `return_dict=True`. If not provided, one will get only the formatted text, optionally tokenized text.
+
+ conversation = [
+ {
+ "role": "user",
+ "content": [
+ {"type": "image", "image": "https://www.ilankelman.org/stopsigns/australia.jpg"},
+ {"type": "text", "text": "Please describe this image in detail."},
+ ],
+ },
+ ]
+
+ Args:
+ conversation (`Union[List[Dict, [str, str]], List[List[Dict[str, str]]]]`):
+ The conversation to format.
+ chat_template (`Optional[str]`, *optional*):
+ The Jinja template to use for formatting the conversation. If not provided, the tokenizer's
+ chat template is used.
+ """
+
+ if chat_template is None:
+ if self.chat_template is not None:
+ chat_template = self.chat_template
+ else:
+ raise ValueError(
+ "No chat template is set for this processor. Please either set the `chat_template` attribute, "
+ "or provide a chat template as an argument. See "
+ "https://huggingface.co/docs/transformers/main/en/chat_templating for more information."
+ )
+
+ # Fill two sets of kwargs that should be used by tokenizer's `apply_chat_template`
+ # and for multimodal data loading. Everything else will be used in `__call__`
+ tokenizer_template_kwargs = {}
+ for tokenizer_key in TokenizerChatTemplateKwargs.__annotations__.keys():
+ default_value = getattr(TokenizerChatTemplateKwargs, tokenizer_key, None)
+ value = kwargs.pop(tokenizer_key, default_value)
+ tokenizer_template_kwargs[tokenizer_key] = value
+
+ mm_load_kwargs = {}
+ for mm_load_key in ChatTemplateLoadKwargs.__annotations__.keys():
+ default_value = getattr(ChatTemplateLoadKwargs, mm_load_key, None)
+ value = kwargs.pop(mm_load_key, default_value)
+ mm_load_kwargs[mm_load_key] = value
+
+ if isinstance(conversation, (list, tuple)) and (
+ isinstance(conversation[0], (list, tuple)) or hasattr(conversation[0], "content")
+ ):
+ is_batched = True
+ conversations = conversation
+ else:
+ is_batched = False
+ conversations = [conversation]
+
+ tokenize = kwargs.pop("tokenize", False)
+ return_dict = kwargs.pop("return_dict", False)
+
+ if tokenize:
+ batch_images, batch_videos = [], []
+ batch_audios = []
+ batch_video_metadata = []
+ for conversation in conversations:
+ images, videos = [], []
+ video_metadata = []
+ for message in conversation:
+ visuals = [content for content in message["content"] if content["type"] in ["image", "video"]]
+ audio_fnames = [
+ content[key]
+ for content in message["content"]
+ for key in ["audio", "url", "path"]
+ if key in content and content["type"] == "audio"
+ ]
+ image_fnames = [
+ vision_info[key]
+ for vision_info in visuals
+ for key in ["image", "url", "path", "base64"]
+ if key in vision_info and vision_info["type"] == "image"
+ ]
+ video_fnames = [
+ vision_info[key]
+ for vision_info in visuals
+ for key in ["video", "url", "path"]
+ if key in vision_info and vision_info["type"] == "video"
+ ]
+
+ for fname in image_fnames:
+ images.append(load_image(fname))
+
+ # Audio models do not accept nested list of audios (yet!) so we construct a flat input audio list
+ if not mm_load_kwargs["load_audio_from_video"]:
+ for fname in audio_fnames:
+ batch_audios.append(load_audio(fname, sampling_rate=mm_load_kwargs["sampling_rate"]))
+ else:
+ for fname in video_fnames:
+ batch_audios.append(load_audio(fname, sampling_rate=mm_load_kwargs["sampling_rate"]))
+
+ for fname in video_fnames:
+ if isinstance(fname, (list, tuple)) and isinstance(fname[0], str):
+ video = [np.array(load_image(image_fname)).T for image_fname in fname]
+ # create a 4D video because `load_video` always returns a 4D array
+ video = np.stack(video)
+ metadata = None
+ logger.warning(
+ "When loading the video from list of images, we cannot infer metadata such as `fps` or `duration`. "
+ "If your model uses this metadata during processing, please load the whole video and let the model sample frames instead."
+ )
+ else:
+ video, metadata = load_video(
+ fname,
+ num_frames=mm_load_kwargs["num_frames"],
+ fps=mm_load_kwargs["video_fps"],
+ backend=mm_load_kwargs["video_load_backend"],
+ sample_indices_fn=mm_load_kwargs["sample_indices_fn"],
+ )
+ videos.append(video)
+ video_metadata.append(metadata)
+
+ # Currently all processors can accept nested list of batches, but not flat list of visuals
+ # So we'll make a batched list of images and let the processor handle it
+ if images:
+ batch_images.append(images)
+ if videos:
+ batch_videos.append(videos)
+ batch_video_metadata.append(video_metadata)
+
+ # Process conversation with video/image information if needed. Then convert into a prompt using Jinja template
+ conversations = self._process_messages_for_chat_template(
+ conversations,
+ batch_images=batch_images,
+ batch_videos=batch_videos,
+ batch_video_metadata=batch_video_metadata,
+ **mm_load_kwargs,
+ )
+
+ prompt = self.tokenizer.apply_chat_template(
+ conversations,
+ chat_template=chat_template,
+ tokenize=False,
+ return_dict=False,
+ **tokenizer_template_kwargs,
+ )
+
+ if not is_batched:
+ prompt = prompt[0]
+
+ if tokenize:
+ # Tokenizer's `apply_chat_template` never adds special tokens when tokenizing
+ # But processor's `apply_chat_template` didn't have an option to tokenize, so users had to format the prompt
+ # and pass it to the processor. Users thus never worried about special tokens relying on processor handling
+ # everything internally. The below line is to keep BC for that and be able to work with model that have
+ # special tokens in the template (consistent with tokenizers). We dont want to raise warning, it will flood command line
+ # without actionable solution for users
+ single_prompt = prompt[0] if is_batched else prompt
+ if self.tokenizer.bos_token is not None and single_prompt.startswith(self.tokenizer.bos_token):
+ kwargs["add_special_tokens"] = False
+
+ out = self(
+ text=prompt,
+ images=batch_images if batch_images else None,
+ videos=batch_videos if batch_videos else None,
+ audio=batch_audios if batch_audios else None,
+ **kwargs,
+ )
+ if return_dict:
+ return out
+ else:
+ return out["input_ids"]
+ return prompt
+
+ def post_process_image_text_to_text(self, generated_outputs, skip_special_tokens=True, **kwargs):
+ """
+ Post-process the output of a vlm to decode the text.
+
+ Args:
+ generated_outputs (`torch.Tensor` or `np.ndarray`):
+ The output of the model `generate` function. The output is expected to be a tensor of shape `(batch_size, sequence_length)`
+ or `(sequence_length,)`.
+ skip_special_tokens (`bool`, *optional*, defaults to `True`):
+ Whether or not to remove special tokens in the output. Argument passed to the tokenizer's `batch_decode` method.
+ **kwargs:
+ Additional arguments to be passed to the tokenizer's `batch_decode method`.
+
+ Returns:
+ `List[str]`: The decoded text.
+ """
+ return self.tokenizer.batch_decode(generated_outputs, skip_special_tokens=skip_special_tokens, **kwargs)
+
+
+def _validate_images_text_input_order(images, text):
+ """
+ For backward compatibility: reverse the order of `images` and `text` inputs if they are swapped.
+ This method should only be called for processors where `images` and `text` have been swapped for uniformization purposes.
+ Note that this method assumes that two `None` inputs are valid inputs. If this is not the case, it should be handled
+ in the processor's `__call__` method before calling this method.
+ """
+
+ def is_url(val) -> bool:
+ return isinstance(val, str) and val.startswith("http")
+
+ def _is_valid_images_input_for_processor(imgs):
+ # If we have an list of images, make sure every image is valid
+ if isinstance(imgs, (list, tuple)):
+ for img in imgs:
+ if not _is_valid_images_input_for_processor(img):
+ return False
+ # If not a list or tuple, we have been given a single image or batched tensor of images
+ elif not (is_valid_image(imgs) or is_url(imgs)):
+ return False
+ return True
+
+ def _is_valid_text_input_for_processor(t):
+ if isinstance(t, str):
+ # Strings are fine
+ return True
+ elif isinstance(t, (list, tuple)):
+ # List are fine as long as they are...
+ if len(t) == 0:
+ # ... not empty
+ return False
+ for t_s in t:
+ return _is_valid_text_input_for_processor(t_s)
+ return False
+
+ def _is_valid(input, validator):
+ return validator(input) or input is None
+
+ images_is_valid = _is_valid(images, _is_valid_images_input_for_processor)
+ images_is_text = _is_valid_text_input_for_processor(images)
+
+ text_is_valid = _is_valid(text, _is_valid_text_input_for_processor)
+ text_is_images = _is_valid_images_input_for_processor(text)
+ # Handle cases where both inputs are valid
+ if images_is_valid and text_is_valid:
+ return images, text
+
+ # Handle cases where inputs need to and can be swapped
+ if (images is None and text_is_images) or (text is None and images_is_text) or (images_is_text and text_is_images):
+ logger.warning_once(
+ "You may have used the wrong order for inputs. `images` should be passed before `text`. "
+ "The `images` and `text` inputs will be swapped. This behavior will be deprecated in transformers v4.47."
+ )
+ return text, images
+
+ raise ValueError("Invalid input type. Check that `images` and/or `text` are valid inputs.")
+
+
+ProcessorMixin.push_to_hub = copy_func(ProcessorMixin.push_to_hub)
+if ProcessorMixin.push_to_hub.__doc__ is not None:
+ ProcessorMixin.push_to_hub.__doc__ = ProcessorMixin.push_to_hub.__doc__.format(
+ object="processor", object_class="AutoProcessor", object_files="processor files"
+ )
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/pytorch_utils.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/pytorch_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..ff3856c9d4573d9b44edeb9a359e8958bf31811b
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/pytorch_utils.py
@@ -0,0 +1,360 @@
+# Copyright 2022 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.
+from __future__ import annotations
+
+import inspect
+from functools import lru_cache, wraps
+from typing import Callable
+
+import torch
+from safetensors.torch import storage_ptr, storage_size
+from torch import nn
+
+from .utils import is_torch_greater_or_equal, is_torch_xla_available, is_torchdynamo_compiling, logging
+
+
+ALL_LAYERNORM_LAYERS = [nn.LayerNorm]
+
+logger = logging.get_logger(__name__)
+
+is_torch_greater_or_equal_than_2_6 = is_torch_greater_or_equal("2.6", accept_dev=True)
+is_torch_greater_or_equal_than_2_4 = is_torch_greater_or_equal("2.4", accept_dev=True)
+is_torch_greater_or_equal_than_2_3 = is_torch_greater_or_equal("2.3", accept_dev=True)
+is_torch_greater_or_equal_than_2_2 = is_torch_greater_or_equal("2.2", accept_dev=True)
+is_torch_greater_or_equal_than_2_1 = is_torch_greater_or_equal("2.1", accept_dev=True)
+
+# For backwards compatibility (e.g. some remote codes on Hub using those variables).
+is_torch_greater_or_equal_than_2_0 = is_torch_greater_or_equal("2.0", accept_dev=True)
+is_torch_greater_or_equal_than_1_13 = is_torch_greater_or_equal("1.13", accept_dev=True)
+is_torch_greater_or_equal_than_1_12 = is_torch_greater_or_equal("1.12", accept_dev=True)
+
+# Cache this result has it's a C FFI call which can be pretty time-consuming
+_torch_distributed_available = torch.distributed.is_available()
+
+if is_torch_greater_or_equal("2.5") and _torch_distributed_available:
+ pass
+
+
+def softmax_backward_data(parent, grad_output, output, dim, self):
+ """
+ A function that calls the internal `_softmax_backward_data` PyTorch method and that adjusts the arguments according
+ to the torch version detected.
+ """
+
+ from torch import _softmax_backward_data
+
+ return _softmax_backward_data(grad_output, output, parent.dim, self.dtype)
+
+
+def prune_linear_layer(layer: nn.Linear, index: torch.LongTensor, dim: int = 0) -> nn.Linear:
+ """
+ Prune a linear layer to keep only entries in index.
+
+ Used to remove heads.
+
+ Args:
+ layer (`torch.nn.Linear`): The layer to prune.
+ index (`torch.LongTensor`): The indices to keep in the layer.
+ dim (`int`, *optional*, defaults to 0): The dimension on which to keep the indices.
+
+ Returns:
+ `torch.nn.Linear`: The pruned layer as a new layer with `requires_grad=True`.
+ """
+ index = index.to(layer.weight.device)
+ W = layer.weight.index_select(dim, index).detach().clone()
+ if layer.bias is not None:
+ if dim == 1:
+ b = layer.bias.detach().clone()
+ else:
+ b = layer.bias[index].detach().clone()
+ new_size = list(layer.weight.size())
+ new_size[dim] = len(index)
+ new_layer = nn.Linear(new_size[1], new_size[0], bias=layer.bias is not None).to(layer.weight.device)
+ new_layer.weight.requires_grad = False
+ new_layer.weight.copy_(W.contiguous())
+ new_layer.weight.requires_grad = True
+ if layer.bias is not None:
+ new_layer.bias.requires_grad = False
+ new_layer.bias.copy_(b.contiguous())
+ new_layer.bias.requires_grad = True
+ return new_layer
+
+
+class Conv1D(nn.Module):
+ """
+ 1D-convolutional layer as defined by Radford et al. for OpenAI GPT (and also used in GPT-2).
+
+ Basically works like a linear layer but the weights are transposed.
+
+ Args:
+ nf (`int`): The number of output features.
+ nx (`int`): The number of input features.
+ """
+
+ def __init__(self, nf, nx):
+ super().__init__()
+ self.nf = nf
+ self.nx = nx
+ self.weight = nn.Parameter(torch.empty(nx, nf))
+ self.bias = nn.Parameter(torch.zeros(nf))
+ nn.init.normal_(self.weight, std=0.02)
+
+ def __repr__(self) -> str:
+ return "Conv1D(nf={nf}, nx={nx})".format(**self.__dict__)
+
+ def forward(self, x):
+ size_out = x.size()[:-1] + (self.nf,)
+ x = torch.addmm(self.bias, x.view(-1, x.size(-1)), self.weight)
+ x = x.view(size_out)
+ return x
+
+
+def prune_conv1d_layer(layer: Conv1D, index: torch.LongTensor, dim: int = 1) -> Conv1D:
+ """
+ Prune a Conv1D layer to keep only entries in index. A Conv1D work as a Linear layer (see e.g. BERT) but the weights
+ are transposed.
+
+ Used to remove heads.
+
+ Args:
+ layer ([`~pytorch_utils.Conv1D`]): The layer to prune.
+ index (`torch.LongTensor`): The indices to keep in the layer.
+ dim (`int`, *optional*, defaults to 1): The dimension on which to keep the indices.
+
+ Returns:
+ [`~pytorch_utils.Conv1D`]: The pruned layer as a new layer with `requires_grad=True`.
+ """
+ index = index.to(layer.weight.device)
+ W = layer.weight.index_select(dim, index).detach().clone()
+ if dim == 0:
+ b = layer.bias.detach().clone()
+ else:
+ b = layer.bias[index].detach().clone()
+ new_size = list(layer.weight.size())
+ new_size[dim] = len(index)
+ new_layer = Conv1D(new_size[1], new_size[0]).to(layer.weight.device)
+ new_layer.weight.requires_grad = False
+ new_layer.weight.copy_(W.contiguous())
+ new_layer.weight.requires_grad = True
+ new_layer.bias.requires_grad = False
+ new_layer.bias.copy_(b.contiguous())
+ new_layer.bias.requires_grad = True
+ return new_layer
+
+
+def prune_layer(layer: nn.Linear | Conv1D, index: torch.LongTensor, dim: int | None = None) -> nn.Linear | Conv1D:
+ """
+ Prune a Conv1D or linear layer to keep only entries in index.
+
+ Used to remove heads.
+
+ Args:
+ layer (`Union[torch.nn.Linear, Conv1D]`): The layer to prune.
+ index (`torch.LongTensor`): The indices to keep in the layer.
+ dim (`int`, *optional*): The dimension on which to keep the indices.
+
+ Returns:
+ `torch.nn.Linear` or [`~pytorch_utils.Conv1D`]: The pruned layer as a new layer with `requires_grad=True`.
+ """
+ if isinstance(layer, nn.Linear):
+ return prune_linear_layer(layer, index, dim=0 if dim is None else dim)
+ elif isinstance(layer, Conv1D):
+ return prune_conv1d_layer(layer, index, dim=1 if dim is None else dim)
+ else:
+ raise ValueError(f"Can't prune layer of class {layer.__class__}")
+
+
+def apply_chunking_to_forward(
+ forward_fn: Callable[..., torch.Tensor],
+ chunk_size: int,
+ chunk_dim: int,
+ *input_tensors,
+) -> torch.Tensor:
+ """
+ This function chunks the `input_tensors` into smaller input tensor parts of size `chunk_size` over the dimension
+ `chunk_dim`. It then applies a layer `forward_fn` to each chunk independently to save memory.
+
+ If the `forward_fn` is independent across the `chunk_dim` this function will yield the same result as directly
+ applying `forward_fn` to `input_tensors`.
+
+ Args:
+ forward_fn (`Callable[..., torch.Tensor]`):
+ The forward function of the model.
+ chunk_size (`int`):
+ The chunk size of a chunked tensor: `num_chunks = len(input_tensors[0]) / chunk_size`.
+ chunk_dim (`int`):
+ The dimension over which the `input_tensors` should be chunked.
+ input_tensors (`Tuple[torch.Tensor]`):
+ The input tensors of `forward_fn` which will be chunked
+
+ Returns:
+ `torch.Tensor`: A tensor with the same shape as the `forward_fn` would have given if applied`.
+
+
+ Examples:
+
+ ```python
+ # rename the usual forward() fn to forward_chunk()
+ def forward_chunk(self, hidden_states):
+ hidden_states = self.decoder(hidden_states)
+ return hidden_states
+
+
+ # implement a chunked forward function
+ def forward(self, hidden_states):
+ return apply_chunking_to_forward(self.forward_chunk, self.chunk_size_lm_head, self.seq_len_dim, hidden_states)
+ ```"""
+
+ assert len(input_tensors) > 0, f"{input_tensors} has to be a tuple/list of tensors"
+
+ # inspect.signature exist since python 3.5 and is a python method -> no problem with backward compatibility
+ num_args_in_forward_chunk_fn = len(inspect.signature(forward_fn).parameters)
+ if num_args_in_forward_chunk_fn != len(input_tensors):
+ raise ValueError(
+ f"forward_chunk_fn expects {num_args_in_forward_chunk_fn} arguments, but only {len(input_tensors)} input "
+ "tensors are given"
+ )
+
+ if chunk_size > 0:
+ tensor_shape = input_tensors[0].shape[chunk_dim]
+ for input_tensor in input_tensors:
+ if input_tensor.shape[chunk_dim] != tensor_shape:
+ raise ValueError(
+ f"All input tenors have to be of the same shape: {tensor_shape}, "
+ f"found shape {input_tensor.shape[chunk_dim]}"
+ )
+
+ if input_tensors[0].shape[chunk_dim] % chunk_size != 0:
+ raise ValueError(
+ f"The dimension to be chunked {input_tensors[0].shape[chunk_dim]} has to be a multiple of the chunk "
+ f"size {chunk_size}"
+ )
+
+ num_chunks = input_tensors[0].shape[chunk_dim] // chunk_size
+
+ # chunk input tensor into tuples
+ input_tensors_chunks = tuple(input_tensor.chunk(num_chunks, dim=chunk_dim) for input_tensor in input_tensors)
+ # apply forward fn to every tuple
+ output_chunks = tuple(forward_fn(*input_tensors_chunk) for input_tensors_chunk in zip(*input_tensors_chunks))
+ # concatenate output at same dimension
+ return torch.cat(output_chunks, dim=chunk_dim)
+
+ return forward_fn(*input_tensors)
+
+
+def find_pruneable_heads_and_indices(
+ heads: list[int], n_heads: int, head_size: int, already_pruned_heads: set[int]
+) -> tuple[set[int], torch.LongTensor]:
+ """
+ Finds the heads and their indices taking `already_pruned_heads` into account.
+
+ Args:
+ heads (`List[int]`): List of the indices of heads to prune.
+ n_heads (`int`): The number of heads in the model.
+ head_size (`int`): The size of each head.
+ already_pruned_heads (`Set[int]`): A set of already pruned heads.
+
+ Returns:
+ `Tuple[Set[int], torch.LongTensor]`: A tuple with the indices of heads to prune taking `already_pruned_heads`
+ into account and the indices of rows/columns to keep in the layer weight.
+ """
+ mask = torch.ones(n_heads, head_size)
+ heads = set(heads) - already_pruned_heads # Convert to set and remove already pruned heads
+ for head in heads:
+ # Compute how many pruned heads are before the head and move the index accordingly
+ head = head - sum(1 if h < head else 0 for h in already_pruned_heads)
+ mask[head] = 0
+ mask = mask.view(-1).contiguous().eq(1)
+ index: torch.LongTensor = torch.arange(len(mask))[mask].long()
+ return heads, index
+
+
+def meshgrid(*tensors: torch.Tensor | list[torch.Tensor], indexing: str | None = None) -> tuple[torch.Tensor, ...]:
+ """
+ Wrapper around torch.meshgrid to avoid warning messages about the introduced `indexing` argument.
+
+ Reference: https://pytorch.org/docs/1.13/generated/torch.meshgrid.html
+ """
+ return torch.meshgrid(*tensors, indexing=indexing)
+
+
+def id_tensor_storage(tensor: torch.Tensor) -> tuple[torch.device, int, int]:
+ """
+ Unique identifier to a tensor storage. Multiple different tensors can share the same underlying storage. For
+ example, "meta" tensors all share the same storage, and thus their identifier will all be equal. This identifier is
+ guaranteed to be unique and constant for this tensor's storage during its lifetime. Two tensor storages with
+ non-overlapping lifetimes may have the same id.
+ """
+ if tensor.device.type == "xla" and is_torch_xla_available():
+ # NOTE: xla tensors dont have storage
+ # use some other unique id to distinguish.
+ # this is a XLA tensor, it must be created using torch_xla's
+ # device. So the following import is safe:
+ import torch_xla
+
+ unique_id = torch_xla._XLAC._xla_get_tensor_id(tensor)
+ else:
+ unique_id = storage_ptr(tensor)
+
+ return tensor.device, unique_id, storage_size(tensor)
+
+
+def isin_mps_friendly(elements: torch.Tensor, test_elements: torch.Tensor | int) -> torch.Tensor:
+ """
+ Same as `torch.isin` without flags, but MPS-friendly. We can remove this function when we stop supporting
+ torch <= 2.3. See https://github.com/pytorch/pytorch/issues/77764#issuecomment-2067838075
+
+ Args:
+ elements (`torch.Tensor`): Input elements
+ test_elements (`torch.Tensor` or `int`): The elements to check against.
+
+ Returns:
+ `torch.Tensor`: A boolean tensor of the same shape as `elements` that is True for `elements` in `test_elements`
+ and False otherwise
+ """
+
+ if elements.device.type == "mps" and not is_torch_greater_or_equal_than_2_4:
+ test_elements = torch.tensor(test_elements)
+ if test_elements.ndim == 0:
+ test_elements = test_elements.unsqueeze(0)
+ return elements.tile(test_elements.shape[0], 1).eq(test_elements.unsqueeze(1)).sum(dim=0).bool().squeeze()
+ else:
+ # Note: don't use named arguments in `torch.isin`, see https://github.com/pytorch/pytorch/issues/126045
+ return torch.isin(elements, test_elements)
+
+
+def compile_compatible_method_lru_cache(*lru_args, **lru_kwargs):
+ """
+ LRU cache decorator from standard functools library, but with a workaround to disable
+ caching when torchdynamo is compiling. Expected to work with class methods.
+ """
+
+ def decorator(func):
+ @wraps(func)
+ def wrapper(self, *args, **kwargs):
+ if not is_torchdynamo_compiling():
+ # Cache the function only if the model is not being compiled
+ # check if the function is already cached, otherwise create it
+ if not hasattr(self, f"_cached_{func.__name__}"):
+ self.__setattr__(
+ f"_cached_{func.__name__}", lru_cache(*lru_args, **lru_kwargs)(func.__get__(self))
+ )
+ return self.__getattribute__(f"_cached_{func.__name__}")(*args, **kwargs)
+ else:
+ # Otherwise, just call the original function
+ return func(self, *args, **kwargs)
+
+ return wrapper
+
+ return decorator
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantization_config.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantization_config.py
new file mode 100644
index 0000000000000000000000000000000000000000..b0f119c58b67b919112884cb2e9253c1c295ebcc
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantization_config.py
@@ -0,0 +1,1835 @@
+#!/usr/bin/env python
+# coding=utf-8
+
+# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
+# Modifications Copyright (C) 2025, Advanced Micro Devices, Inc. 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 copy
+import dataclasses
+import importlib.metadata
+import json
+import os
+from dataclasses import dataclass
+from enum import Enum
+from inspect import Parameter, signature
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+from packaging import version
+
+from ..utils import (
+ is_auto_awq_available,
+ is_compressed_tensors_available,
+ is_gptqmodel_available,
+ is_hqq_available,
+ is_quark_available,
+ is_torch_available,
+ is_torchao_available,
+ logging,
+)
+from .import_utils import is_auto_gptq_available
+
+
+if is_torch_available():
+ import torch
+
+logger = logging.get_logger(__name__)
+
+
+class QuantizationMethod(str, Enum):
+ BITS_AND_BYTES = "bitsandbytes"
+ GPTQ = "gptq"
+ AWQ = "awq"
+ AQLM = "aqlm"
+ VPTQ = "vptq"
+ QUANTO = "quanto"
+ EETQ = "eetq"
+ HIGGS = "higgs"
+ HQQ = "hqq"
+ COMPRESSED_TENSORS = "compressed-tensors"
+ FBGEMM_FP8 = "fbgemm_fp8"
+ TORCHAO = "torchao"
+ BITNET = "bitnet"
+ SPQR = "spqr"
+ FP8 = "fp8"
+ QUARK = "quark"
+
+
+class AWQLinearVersion(str, Enum):
+ GEMM = "gemm"
+ GEMV = "gemv"
+ EXLLAMA = "exllama"
+ IPEX = "ipex"
+
+ @staticmethod
+ def from_str(version: str):
+ version = version.lower()
+ if version == "gemm":
+ return AWQLinearVersion.GEMM
+ elif version == "gemv":
+ return AWQLinearVersion.GEMV
+ elif version == "exllama":
+ return AWQLinearVersion.EXLLAMA
+ elif version == "ipex":
+ return AWQLinearVersion.IPEX
+ else:
+ raise ValueError(f"Unknown AWQLinearVersion {version}")
+
+
+class AwqBackendPackingMethod(str, Enum):
+ AUTOAWQ = "autoawq"
+ LLMAWQ = "llm-awq"
+
+
+@dataclass
+class QuantizationConfigMixin:
+ """
+ Mixin class for quantization config
+ """
+
+ quant_method: QuantizationMethod
+
+ @classmethod
+ def from_dict(cls, config_dict, return_unused_kwargs=False, **kwargs):
+ """
+ Instantiates a [`QuantizationConfigMixin`] from a Python dictionary of parameters.
+
+ Args:
+ config_dict (`Dict[str, Any]`):
+ Dictionary that will be used to instantiate the configuration object.
+ return_unused_kwargs (`bool`,*optional*, defaults to `False`):
+ Whether or not to return a list of unused keyword arguments. Used for `from_pretrained` method in
+ `PreTrainedModel`.
+ kwargs (`Dict[str, Any]`):
+ Additional parameters from which to initialize the configuration object.
+
+ Returns:
+ [`QuantizationConfigMixin`]: The configuration object instantiated from those parameters.
+ """
+ config = cls(**config_dict)
+
+ to_remove = []
+ for key, value in kwargs.items():
+ if hasattr(config, key):
+ setattr(config, key, value)
+ to_remove.append(key)
+ for key in to_remove:
+ kwargs.pop(key, None)
+
+ if return_unused_kwargs:
+ return config, kwargs
+ else:
+ return config
+
+ def to_json_file(self, json_file_path: Union[str, os.PathLike]):
+ """
+ Save this instance to a JSON file.
+
+ Args:
+ json_file_path (`str` or `os.PathLike`):
+ Path to the JSON file in which this configuration instance's parameters will be saved.
+ use_diff (`bool`, *optional*, defaults to `True`):
+ If set to `True`, only the difference between the config instance and the default
+ `QuantizationConfig()` is serialized to JSON file.
+ """
+ with open(json_file_path, "w", encoding="utf-8") as writer:
+ config_dict = self.to_dict()
+ json_string = json.dumps(config_dict, indent=2, sort_keys=True) + "\n"
+
+ writer.write(json_string)
+
+ def to_dict(self) -> Dict[str, Any]:
+ """
+ Serializes this instance to a Python dictionary. Returns:
+ `Dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance.
+ """
+ return copy.deepcopy(self.__dict__)
+
+ def __iter__(self):
+ """allows `dict(obj)` for situations where obj may be a dict or QuantizationConfigMixin"""
+ for attr, value in copy.deepcopy(self.__dict__).items():
+ yield attr, value
+
+ def __repr__(self):
+ return f"{self.__class__.__name__} {self.to_json_string()}"
+
+ def to_json_string(self, use_diff: bool = True) -> str:
+ """
+ Serializes this instance to a JSON string.
+
+ Args:
+ use_diff (`bool`, *optional*, defaults to `True`):
+ If set to `True`, only the difference between the config instance and the default `PretrainedConfig()`
+ is serialized to JSON string.
+
+ Returns:
+ `str`: String containing all the attributes that make up this configuration instance in JSON format.
+ """
+ if use_diff is True:
+ config_dict = self.to_diff_dict()
+ else:
+ config_dict = self.to_dict()
+ return json.dumps(config_dict, indent=2, sort_keys=True) + "\n"
+
+ def update(self, **kwargs):
+ """
+ Updates attributes of this class instance with attributes from `kwargs` if they match existing attributes,
+ returning all the unused kwargs.
+
+ Args:
+ kwargs (`Dict[str, Any]`):
+ Dictionary of attributes to tentatively update this class.
+
+ Returns:
+ `Dict[str, Any]`: Dictionary containing all the key-value pairs that were not used to update the instance.
+ """
+ to_remove = []
+ for key, value in kwargs.items():
+ if hasattr(self, key):
+ setattr(self, key, value)
+ to_remove.append(key)
+
+ # Remove all the attributes that were updated, without modifying the input dict
+ unused_kwargs = {key: value for key, value in kwargs.items() if key not in to_remove}
+ return unused_kwargs
+
+
+@dataclass
+class HqqConfig(QuantizationConfigMixin):
+ """
+ This is wrapper around hqq's BaseQuantizeConfig.
+
+ Args:
+ nbits (`int`, *optional*, defaults to 4):
+ Number of bits. Supported values are (8, 4, 3, 2, 1).
+ group_size (`int`, *optional*, defaults to 64):
+ Group-size value. Supported values are any value that is divisble by weight.shape[axis]).
+ view_as_float (`bool`, *optional*, defaults to `False`):
+ View the quantized weight as float (used in distributed training) if set to `True`.
+ axis (`Optional[int]`, *optional*):
+ Axis along which grouping is performed. Supported values are 0 or 1.
+ dynamic_config (dict, *optional*):
+ Parameters for dynamic configuration. The key is the name tag of the layer and the value is a quantization config.
+ If set, each layer specified by its id will use its dedicated quantization configuration.
+ skip_modules (`List[str]`, *optional*, defaults to `['lm_head']`):
+ List of `nn.Linear` layers to skip.
+ kwargs (`Dict[str, Any]`, *optional*):
+ Additional parameters from which to initialize the configuration object.
+ """
+
+ def __init__(
+ self,
+ nbits: int = 4,
+ group_size: int = 64,
+ view_as_float: bool = False,
+ axis: Optional[int] = None,
+ dynamic_config: Optional[dict] = None,
+ skip_modules: List[str] = ["lm_head"],
+ **kwargs,
+ ):
+ if is_hqq_available():
+ from hqq.core.quantize import BaseQuantizeConfig as HQQBaseQuantizeConfig
+ else:
+ raise ImportError(
+ "A valid HQQ version (>=0.2.1) is not available. Please follow the instructions to install it: `https://github.com/mobiusml/hqq/`."
+ )
+
+ for deprecated_key in ["quant_zero", "quant_scale", "offload_meta"]:
+ if deprecated_key in kwargs:
+ logger.info(
+ deprecated_key + " is deprecated. This parameter will be ignored in quantization settings."
+ )
+
+ if axis is None:
+ axis = 1
+ logger.info("Setting axis=1 as faster backends such as TorchAO or BitBlas are only compatible with it.")
+
+ if axis not in [0, 1]:
+ raise ValueError("Invalid axis value. Only 0 and 1 are allowed.")
+
+ if dynamic_config is not None:
+ self.quant_config = {}
+ for key in dynamic_config:
+ self.quant_config[key] = HQQBaseQuantizeConfig(**dynamic_config[key])
+ else:
+ self.quant_config = HQQBaseQuantizeConfig(
+ **{
+ "nbits": nbits,
+ "group_size": group_size,
+ "view_as_float": view_as_float,
+ "axis": axis,
+ }
+ )
+
+ self.quant_method = QuantizationMethod.HQQ
+ self.skip_modules = skip_modules
+
+ self.post_init()
+
+ def post_init(self):
+ r"""
+ Safety checker that arguments are correct - also replaces some NoneType arguments with their default values.
+ """
+ pass
+
+ @classmethod
+ def from_dict(cls, config: Dict[str, Any]):
+ """
+ Override from_dict, used in AutoQuantizationConfig.from_dict in quantizers/auto.py
+ """
+ instance = cls()
+ instance.quant_config = config["quant_config"]
+ instance.skip_modules = config["skip_modules"]
+ return instance
+
+ def to_dict(self) -> Dict[str, Any]:
+ """
+ Serializes this instance to a Python dictionary. Returns:
+ `Dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance.
+ """
+ return {
+ "quant_config": self.quant_config,
+ "quant_method": self.quant_method,
+ "skip_modules": self.skip_modules,
+ }
+
+ def __repr__(self):
+ config_dict = self.to_dict()
+ return f"{self.__class__.__name__} {json.dumps(config_dict, indent=2, sort_keys=True)}\n"
+
+ def to_diff_dict(self) -> Dict[str, Any]:
+ """
+ Removes all attributes from config which correspond to the default config attributes for better readability and
+ serializes to a Python dictionary.
+ Returns:
+ `Dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance,
+ """
+ config_dict = self.to_dict()
+
+ # get the default config dict
+ default_config_dict = HqqConfig().to_dict()
+
+ serializable_config_dict = {}
+
+ # only serialize values that differ from the default config
+ for key, value in config_dict.items():
+ if value != default_config_dict[key]:
+ serializable_config_dict[key] = value
+
+ return serializable_config_dict
+
+
+@dataclass
+class BitsAndBytesConfig(QuantizationConfigMixin):
+ """
+ This is a wrapper class about all possible attributes and features that you can play with a model that has been
+ loaded using `bitsandbytes`.
+
+ This replaces `load_in_8bit` or `load_in_4bit`therefore both options are mutually exclusive.
+
+ Currently only supports `LLM.int8()`, `FP4`, and `NF4` quantization. If more methods are added to `bitsandbytes`,
+ then more arguments will be added to this class.
+
+ Args:
+ load_in_8bit (`bool`, *optional*, defaults to `False`):
+ This flag is used to enable 8-bit quantization with LLM.int8().
+ load_in_4bit (`bool`, *optional*, defaults to `False`):
+ This flag is used to enable 4-bit quantization by replacing the Linear layers with FP4/NF4 layers from
+ `bitsandbytes`.
+ llm_int8_threshold (`float`, *optional*, defaults to 6.0):
+ This corresponds to the outlier threshold for outlier detection as described in `LLM.int8() : 8-bit Matrix
+ Multiplication for Transformers at Scale` paper: https://arxiv.org/abs/2208.07339 Any hidden states value
+ that is above this threshold will be considered an outlier and the operation on those values will be done
+ in fp16. Values are usually normally distributed, that is, most values are in the range [-3.5, 3.5], but
+ there are some exceptional systematic outliers that are very differently distributed for large models.
+ These outliers are often in the interval [-60, -6] or [6, 60]. Int8 quantization works well for values of
+ magnitude ~5, but beyond that, there is a significant performance penalty. A good default threshold is 6,
+ but a lower threshold might be needed for more unstable models (small models, fine-tuning).
+ llm_int8_skip_modules (`List[str]`, *optional*):
+ An explicit list of the modules that we do not want to convert in 8-bit. This is useful for models such as
+ Jukebox that has several heads in different places and not necessarily at the last position. For example
+ for `CausalLM` models, the last `lm_head` is kept in its original `dtype`.
+ llm_int8_enable_fp32_cpu_offload (`bool`, *optional*, defaults to `False`):
+ This flag is used for advanced use cases and users that are aware of this feature. If you want to split
+ your model in different parts and run some parts in int8 on GPU and some parts in fp32 on CPU, you can use
+ this flag. This is useful for offloading large models such as `google/flan-t5-xxl`. Note that the int8
+ operations will not be run on CPU.
+ llm_int8_has_fp16_weight (`bool`, *optional*, defaults to `False`):
+ This flag runs LLM.int8() with 16-bit main weights. This is useful for fine-tuning as the weights do not
+ have to be converted back and forth for the backward pass.
+ bnb_4bit_compute_dtype (`torch.dtype` or str, *optional*, defaults to `torch.float32`):
+ This sets the computational type which might be different than the input type. For example, inputs might be
+ fp32, but computation can be set to bf16 for speedups.
+ bnb_4bit_quant_type (`str`, *optional*, defaults to `"fp4"`):
+ This sets the quantization data type in the bnb.nn.Linear4Bit layers. Options are FP4 and NF4 data types
+ which are specified by `fp4` or `nf4`.
+ bnb_4bit_use_double_quant (`bool`, *optional*, defaults to `False`):
+ This flag is used for nested quantization where the quantization constants from the first quantization are
+ quantized again.
+ bnb_4bit_quant_storage (`torch.dtype` or str, *optional*, defaults to `torch.uint8`):
+ This sets the storage type to pack the quanitzed 4-bit prarams.
+ kwargs (`Dict[str, Any]`, *optional*):
+ Additional parameters from which to initialize the configuration object.
+ """
+
+ def __init__(
+ self,
+ load_in_8bit=False,
+ load_in_4bit=False,
+ llm_int8_threshold=6.0,
+ llm_int8_skip_modules=None,
+ llm_int8_enable_fp32_cpu_offload=False,
+ llm_int8_has_fp16_weight=False,
+ bnb_4bit_compute_dtype=None,
+ bnb_4bit_quant_type="fp4",
+ bnb_4bit_use_double_quant=False,
+ bnb_4bit_quant_storage=None,
+ **kwargs,
+ ):
+ self.quant_method = QuantizationMethod.BITS_AND_BYTES
+
+ if load_in_4bit and load_in_8bit:
+ raise ValueError("load_in_4bit and load_in_8bit are both True, but only one can be used at the same time")
+
+ self._load_in_8bit = load_in_8bit
+ self._load_in_4bit = load_in_4bit
+ self.llm_int8_threshold = llm_int8_threshold
+ self.llm_int8_skip_modules = llm_int8_skip_modules
+ self.llm_int8_enable_fp32_cpu_offload = llm_int8_enable_fp32_cpu_offload
+ self.llm_int8_has_fp16_weight = llm_int8_has_fp16_weight
+ self.bnb_4bit_quant_type = bnb_4bit_quant_type
+ self.bnb_4bit_use_double_quant = bnb_4bit_use_double_quant
+
+ if bnb_4bit_compute_dtype is None:
+ self.bnb_4bit_compute_dtype = torch.float32
+ elif isinstance(bnb_4bit_compute_dtype, str):
+ self.bnb_4bit_compute_dtype = getattr(torch, bnb_4bit_compute_dtype)
+ elif isinstance(bnb_4bit_compute_dtype, torch.dtype):
+ self.bnb_4bit_compute_dtype = bnb_4bit_compute_dtype
+ else:
+ raise ValueError("bnb_4bit_compute_dtype must be a string or a torch.dtype")
+
+ if bnb_4bit_quant_storage is None:
+ self.bnb_4bit_quant_storage = torch.uint8
+ elif isinstance(bnb_4bit_quant_storage, str):
+ if bnb_4bit_quant_storage not in ["float16", "float32", "int8", "uint8", "float64", "bfloat16"]:
+ raise ValueError(
+ "`bnb_4bit_quant_storage` must be a valid string (one of 'float16', 'float32', 'int8', 'uint8', 'float64', 'bfloat16') "
+ )
+ self.bnb_4bit_quant_storage = getattr(torch, bnb_4bit_quant_storage)
+ elif isinstance(bnb_4bit_quant_storage, torch.dtype):
+ self.bnb_4bit_quant_storage = bnb_4bit_quant_storage
+ else:
+ raise ValueError("bnb_4bit_quant_storage must be a string or a torch.dtype")
+
+ if kwargs:
+ logger.info(f"Unused kwargs: {list(kwargs.keys())}. These kwargs are not used in {self.__class__}.")
+
+ self.post_init()
+
+ @property
+ def load_in_4bit(self):
+ return self._load_in_4bit
+
+ @load_in_4bit.setter
+ def load_in_4bit(self, value: bool):
+ if not isinstance(value, bool):
+ raise TypeError("load_in_4bit must be a boolean")
+
+ if self.load_in_8bit and value:
+ raise ValueError("load_in_4bit and load_in_8bit are both True, but only one can be used at the same time")
+ self._load_in_4bit = value
+
+ @property
+ def load_in_8bit(self):
+ return self._load_in_8bit
+
+ @load_in_8bit.setter
+ def load_in_8bit(self, value: bool):
+ if not isinstance(value, bool):
+ raise TypeError("load_in_8bit must be a boolean")
+
+ if self.load_in_4bit and value:
+ raise ValueError("load_in_4bit and load_in_8bit are both True, but only one can be used at the same time")
+ self._load_in_8bit = value
+
+ def post_init(self):
+ r"""
+ Safety checker that arguments are correct - also replaces some NoneType arguments with their default values.
+ """
+ if not isinstance(self.load_in_4bit, bool):
+ raise TypeError("load_in_4bit must be a boolean")
+
+ if not isinstance(self.load_in_8bit, bool):
+ raise TypeError("load_in_8bit must be a boolean")
+
+ if not isinstance(self.llm_int8_threshold, float):
+ raise TypeError("llm_int8_threshold must be a float")
+
+ if self.llm_int8_skip_modules is not None and not isinstance(self.llm_int8_skip_modules, list):
+ raise TypeError("llm_int8_skip_modules must be a list of strings")
+ if not isinstance(self.llm_int8_enable_fp32_cpu_offload, bool):
+ raise TypeError("llm_int8_enable_fp32_cpu_offload must be a boolean")
+
+ if not isinstance(self.llm_int8_has_fp16_weight, bool):
+ raise TypeError("llm_int8_has_fp16_weight must be a boolean")
+
+ if self.bnb_4bit_compute_dtype is not None and not isinstance(self.bnb_4bit_compute_dtype, torch.dtype):
+ raise TypeError("bnb_4bit_compute_dtype must be torch.dtype")
+
+ if not isinstance(self.bnb_4bit_quant_type, str):
+ raise TypeError("bnb_4bit_quant_type must be a string")
+
+ if not isinstance(self.bnb_4bit_use_double_quant, bool):
+ raise TypeError("bnb_4bit_use_double_quant must be a boolean")
+
+ if self.load_in_4bit and not version.parse(importlib.metadata.version("bitsandbytes")) >= version.parse(
+ "0.39.0"
+ ):
+ raise ValueError(
+ "4 bit quantization requires bitsandbytes>=0.39.0 - please upgrade your bitsandbytes version"
+ )
+
+ def is_quantizable(self):
+ r"""
+ Returns `True` if the model is quantizable, `False` otherwise.
+ """
+ return self.load_in_8bit or self.load_in_4bit
+
+ def quantization_method(self):
+ r"""
+ This method returns the quantization method used for the model. If the model is not quantizable, it returns
+ `None`.
+ """
+ if self.load_in_8bit:
+ return "llm_int8"
+ elif self.load_in_4bit and self.bnb_4bit_quant_type == "fp4":
+ return "fp4"
+ elif self.load_in_4bit and self.bnb_4bit_quant_type == "nf4":
+ return "nf4"
+ else:
+ return None
+
+ def to_dict(self) -> Dict[str, Any]:
+ """
+ Serializes this instance to a Python dictionary. Returns:
+ `Dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance.
+ """
+ output = copy.deepcopy(self.__dict__)
+ output["bnb_4bit_compute_dtype"] = str(output["bnb_4bit_compute_dtype"]).split(".")[1]
+ output["bnb_4bit_quant_storage"] = str(output["bnb_4bit_quant_storage"]).split(".")[1]
+ output["load_in_4bit"] = self.load_in_4bit
+ output["load_in_8bit"] = self.load_in_8bit
+
+ return output
+
+ def __repr__(self):
+ config_dict = self.to_dict()
+ return f"{self.__class__.__name__} {json.dumps(config_dict, indent=2, sort_keys=True)}\n"
+
+ def to_diff_dict(self) -> Dict[str, Any]:
+ """
+ Removes all attributes from config which correspond to the default config attributes for better readability and
+ serializes to a Python dictionary.
+
+ Returns:
+ `Dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance,
+ """
+ config_dict = self.to_dict()
+
+ # get the default config dict
+ default_config_dict = BitsAndBytesConfig().to_dict()
+
+ serializable_config_dict = {}
+
+ # only serialize values that differ from the default config
+ for key, value in config_dict.items():
+ if value != default_config_dict[key]:
+ serializable_config_dict[key] = value
+
+ return serializable_config_dict
+
+
+class ExllamaVersion(int, Enum):
+ ONE = 1
+ TWO = 2
+
+
+@dataclass
+class GPTQConfig(QuantizationConfigMixin):
+ """
+ This is a wrapper class about all possible attributes and features that you can play with a model that has been
+ loaded using `optimum` api for gptq quantization relying on auto_gptq backend.
+
+ Args:
+ bits (`int`):
+ The number of bits to quantize to, supported numbers are (2, 3, 4, 8).
+ tokenizer (`str` or `PreTrainedTokenizerBase`, *optional*):
+ The tokenizer used to process the dataset. You can pass either:
+ - A custom tokenizer object.
+ - A string, the *model id* of a predefined tokenizer hosted inside a model repo on huggingface.co.
+ - A path to a *directory* containing vocabulary files required by the tokenizer, for instance saved
+ using the [`~PreTrainedTokenizer.save_pretrained`] method, e.g., `./my_model_directory/`.
+ dataset (`Union[List[str]]`, *optional*):
+ The dataset used for quantization. You can provide your own dataset in a list of string or just use the
+ original datasets used in GPTQ paper ['wikitext2','c4','c4-new']
+ group_size (`int`, *optional*, defaults to 128):
+ The group size to use for quantization. Recommended value is 128 and -1 uses per-column quantization.
+ damp_percent (`float`, *optional*, defaults to 0.1):
+ The percent of the average Hessian diagonal to use for dampening. Recommended value is 0.1.
+ desc_act (`bool`, *optional*, defaults to `False`):
+ Whether to quantize columns in order of decreasing activation size. Setting it to False can significantly
+ speed up inference but the perplexity may become slightly worse. Also known as act-order.
+ sym (`bool`, *optional*, defaults to `True`):
+ Whether to use symetric quantization.
+ true_sequential (`bool`, *optional*, defaults to `True`):
+ Whether to perform sequential quantization even within a single Transformer block. Instead of quantizing
+ the entire block at once, we perform layer-wise quantization. As a result, each layer undergoes
+ quantization using inputs that have passed through the previously quantized layers.
+ checkpoint_format (`str`, *optional*, defaults to `"gptq"`):
+ GPTQ weight format. `gptq`(v1) is supported by both gptqmodel and auto-gptq. `gptq_v2` is gptqmodel only.
+ meta (`Dict[str, any]`, *optional*):
+ Properties, such as tooling:version, that do not directly contributes to quantization or quant inference are stored in meta.
+ i.e. `meta.quantizer`: ["optimum:_version_", "gptqmodel:_version_"]
+ backend (`str`, *optional*):
+ Controls which gptq kernel to be used. Valid values for gptqmodel are `auto`, `auto_trainable` and more. For auto-gptq, only
+ valid value is None and `auto_trainable`. Ref gptqmodel backends: https://github.com/ModelCloud/GPTQModel/blob/main/gptqmodel/utils/backend.py
+ use_cuda_fp16 (`bool`, *optional*, defaults to `False`):
+ Whether or not to use optimized cuda kernel for fp16 model. Need to have model in fp16. Auto-gptq only.
+ model_seqlen (`int`, *optional*):
+ The maximum sequence length that the model can take.
+ block_name_to_quantize (`str`, *optional*):
+ The transformers block name to quantize. If None, we will infer the block name using common patterns (e.g. model.layers)
+ module_name_preceding_first_block (`List[str]`, *optional*):
+ The layers that are preceding the first Transformer block.
+ batch_size (`int`, *optional*, defaults to 1):
+ The batch size used when processing the dataset
+ pad_token_id (`int`, *optional*):
+ The pad token id. Needed to prepare the dataset when `batch_size` > 1.
+ use_exllama (`bool`, *optional*):
+ Whether to use exllama backend. Defaults to `True` if unset. Only works with `bits` = 4.
+ max_input_length (`int`, *optional*):
+ The maximum input length. This is needed to initialize a buffer that depends on the maximum expected input
+ length. It is specific to the exllama backend with act-order.
+ exllama_config (`Dict[str, Any]`, *optional*):
+ The exllama config. You can specify the version of the exllama kernel through the `version` key. Defaults
+ to `{"version": 1}` if unset.
+ cache_block_outputs (`bool`, *optional*, defaults to `True`):
+ Whether to cache block outputs to reuse as inputs for the succeeding block.
+ modules_in_block_to_quantize (`List[List[str]]`, *optional*):
+ List of list of module names to quantize in the specified block. This argument is useful to exclude certain linear modules from being quantized.
+ The block to quantize can be specified by setting `block_name_to_quantize`. We will quantize each list sequentially. If not set, we will quantize all linear layers.
+ Example: `modules_in_block_to_quantize =[["self_attn.k_proj", "self_attn.v_proj", "self_attn.q_proj"], ["self_attn.o_proj"]]`.
+ In this example, we will first quantize the q,k,v layers simultaneously since they are independent.
+ Then, we will quantize `self_attn.o_proj` layer with the q,k,v layers quantized. This way, we will get
+ better results since it reflects the real input `self_attn.o_proj` will get when the model is quantized.
+ """
+
+ def __init__(
+ self,
+ bits: int,
+ tokenizer: Any = None,
+ dataset: Optional[Union[List[str], str]] = None,
+ group_size: int = 128,
+ damp_percent: float = 0.1,
+ desc_act: bool = False,
+ sym: bool = True,
+ true_sequential: bool = True,
+ checkpoint_format: str = "gptq",
+ meta: Optional[Dict[str, Any]] = None,
+ backend: Optional[str] = None,
+ use_cuda_fp16: bool = False,
+ model_seqlen: Optional[int] = None,
+ block_name_to_quantize: Optional[str] = None,
+ module_name_preceding_first_block: Optional[List[str]] = None,
+ batch_size: int = 1,
+ pad_token_id: Optional[int] = None,
+ use_exllama: Optional[bool] = None,
+ max_input_length: Optional[int] = None,
+ exllama_config: Optional[Dict[str, Any]] = None,
+ cache_block_outputs: bool = True,
+ modules_in_block_to_quantize: Optional[List[List[str]]] = None,
+ **kwargs,
+ ):
+ self.quant_method = QuantizationMethod.GPTQ
+ self.bits = bits
+ self.tokenizer = tokenizer
+ self.dataset = dataset
+ self.group_size = group_size
+ self.damp_percent = damp_percent
+ self.desc_act = desc_act
+ self.sym = sym
+ self.true_sequential = true_sequential
+ self.checkpoint_format = checkpoint_format.lower()
+ self.meta = meta
+ self.backend = backend.lower() if isinstance(backend, str) else backend
+ self.use_cuda_fp16 = use_cuda_fp16
+ self.model_seqlen = model_seqlen
+ self.block_name_to_quantize = block_name_to_quantize
+ self.module_name_preceding_first_block = module_name_preceding_first_block
+ self.batch_size = batch_size
+ self.pad_token_id = pad_token_id
+ self.use_exllama = use_exllama
+ self.max_input_length = max_input_length
+ self.exllama_config = exllama_config
+ self.cache_block_outputs = cache_block_outputs
+ self.modules_in_block_to_quantize = modules_in_block_to_quantize
+ self.post_init()
+
+ def get_loading_attributes(self):
+ attibutes_dict = copy.deepcopy(self.__dict__)
+ loading_attibutes = [
+ "use_exllama",
+ "exllama_config",
+ "use_cuda_fp16",
+ "max_input_length",
+ "backend",
+ ]
+ loading_attibutes_dict = {i: j for i, j in attibutes_dict.items() if i in loading_attibutes}
+ return loading_attibutes_dict
+
+ def post_init(self):
+ r"""
+ Safety checker that arguments are correct
+ """
+ if self.bits not in [2, 3, 4, 8]:
+ raise ValueError(f"Only support quantization to [2,3,4,8] bits but found {self.bits}")
+ if self.group_size != -1 and self.group_size <= 0:
+ raise ValueError("group_size must be greater than 0 or equal to -1")
+ if not (0 < self.damp_percent < 1):
+ raise ValueError("damp_percent must between 0 and 1.")
+ if self.dataset is not None:
+ if isinstance(self.dataset, str):
+ if self.dataset in ["ptb", "ptb-new"]:
+ raise ValueError(
+ f"""{self.dataset} dataset was deprecated. You can only choose between
+ ['wikitext2','c4','c4-new']"""
+ )
+ if self.dataset not in ["wikitext2", "c4", "c4-new"]:
+ raise ValueError(
+ f"""You have entered a string value for dataset. You can only choose between
+ ['wikitext2','c4','c4-new'], but we found {self.dataset}"""
+ )
+ elif not isinstance(self.dataset, list):
+ raise ValueError(
+ f"""dataset needs to be either a list of string or a value in
+ ['wikitext2','c4','c4-new'], but we found {self.dataset}"""
+ )
+
+ # make sure backend is back/forward compatible with both gptqmodel (full) and auto-gptq (partial)
+ if is_gptqmodel_available():
+ # convert auto-gptq control into gptqmodel backend
+ if self.backend is None:
+ self.backend = "auto_trainable" if self.use_exllama is not None and not self.use_exllama else "auto"
+ else:
+ # convert gptqmodel backend `auto_trainable` into auto-gptq control
+ if self.backend == "auto_trainable":
+ self.use_exllama = False
+
+ # auto-gptq specific kernel control logic
+ if self.use_exllama is None:
+ # New default behaviour
+ self.use_exllama = True
+
+ if self.exllama_config is None:
+ self.exllama_config = {"version": ExllamaVersion.ONE}
+ else:
+ if "version" not in self.exllama_config:
+ raise ValueError("`exllama_config` needs to have a `version` key.")
+ elif self.exllama_config["version"] not in [ExllamaVersion.ONE, ExllamaVersion.TWO]:
+ exllama_version = self.exllama_config["version"]
+ raise ValueError(
+ f"Only supported versions are in [ExllamaVersion.ONE, ExllamaVersion.TWO] - not recognized version {exllama_version}"
+ )
+
+ if self.bits == 4 and self.use_exllama:
+ if self.exllama_config["version"] == ExllamaVersion.ONE:
+ logger.info(
+ "You have activated exllama backend. Note that you can get better inference "
+ "speed using exllamav2 kernel by setting `exllama_config`."
+ )
+ elif self.exllama_config["version"] == ExllamaVersion.TWO:
+ if is_auto_gptq_available():
+ optimum_version = version.parse(importlib.metadata.version("optimum"))
+ autogptq_version = version.parse(importlib.metadata.version("auto_gptq"))
+ if optimum_version <= version.parse("1.13.2") or autogptq_version <= version.parse("0.4.2"):
+ raise ValueError(
+ f"You need optimum > 1.13.2 and auto-gptq > 0.4.2 . Make sure to have that version installed - detected version : optimum {optimum_version} and autogptq {autogptq_version}"
+ )
+ if self.modules_in_block_to_quantize is not None:
+ optimum_version = version.parse(importlib.metadata.version("optimum"))
+ if optimum_version < version.parse("1.15.0"):
+ raise ValueError(
+ "You current version of `optimum` does not support `modules_in_block_to_quantize` quantization argument, please upgrade `optimum` package to a version superior than 1.15.0 ."
+ )
+
+ def to_dict(self):
+ config_dict = super().to_dict()
+ config_dict.pop("disable_exllama", None)
+ return config_dict
+
+ def to_dict_optimum(self):
+ """
+ Get compatible dict for optimum gptq config
+ """
+ quant_dict = self.to_dict()
+ # make it compatible with optimum config
+ quant_dict["disable_exllama"] = not self.use_exllama
+ return quant_dict
+
+ @classmethod
+ def from_dict_optimum(cls, config_dict):
+ """
+ Get compatible class with optimum gptq config dict
+ """
+
+ if "disable_exllama" in config_dict:
+ config_dict["use_exllama"] = not config_dict["disable_exllama"]
+ # switch to None to not trigger the warning
+ config_dict.pop("disable_exllama")
+
+ config = cls(**config_dict)
+ return config
+
+
+@dataclass
+class AwqConfig(QuantizationConfigMixin):
+ """
+ This is a wrapper class about all possible attributes and features that you can play with a model that has been
+ loaded using `auto-awq` library awq quantization relying on auto_awq backend.
+
+ Args:
+ bits (`int`, *optional*, defaults to 4):
+ The number of bits to quantize to.
+ group_size (`int`, *optional*, defaults to 128):
+ The group size to use for quantization. Recommended value is 128 and -1 uses per-column quantization.
+ zero_point (`bool`, *optional*, defaults to `True`):
+ Whether to use zero point quantization.
+ version (`AWQLinearVersion`, *optional*, defaults to `AWQLinearVersion.GEMM`):
+ The version of the quantization algorithm to use. GEMM is better for big batch_size (e.g. >= 8) otherwise,
+ GEMV is better (e.g. < 8 ). GEMM models are compatible with Exllama kernels.
+ backend (`AwqBackendPackingMethod`, *optional*, defaults to `AwqBackendPackingMethod.AUTOAWQ`):
+ The quantization backend. Some models might be quantized using `llm-awq` backend. This is useful for users
+ that quantize their own models using `llm-awq` library.
+ do_fuse (`bool`, *optional*, defaults to `False`):
+ Whether to fuse attention and mlp layers together for faster inference
+ fuse_max_seq_len (`int`, *optional*):
+ The Maximum sequence length to generate when using fusing.
+ modules_to_fuse (`dict`, *optional*, default to `None`):
+ Overwrite the natively supported fusing scheme with the one specified by the users.
+ modules_to_not_convert (`list`, *optional*, default to `None`):
+ The list of modules to not quantize, useful for quantizing models that explicitly require to have
+ some modules left in their original precision (e.g. Whisper encoder, Llava encoder, Mixtral gate layers).
+ Note you cannot quantize directly with transformers, please refer to `AutoAWQ` documentation for quantizing HF models.
+ exllama_config (`Dict[str, Any]`, *optional*):
+ You can specify the version of the exllama kernel through the `version` key, the maximum sequence
+ length through the `max_input_len` key, and the maximum batch size through the `max_batch_size` key.
+ Defaults to `{"version": 2, "max_input_len": 2048, "max_batch_size": 8}` if unset.
+ """
+
+ def __init__(
+ self,
+ bits: int = 4,
+ group_size: int = 128,
+ zero_point: bool = True,
+ version: AWQLinearVersion = AWQLinearVersion.GEMM,
+ backend: AwqBackendPackingMethod = AwqBackendPackingMethod.AUTOAWQ,
+ do_fuse: Optional[bool] = None,
+ fuse_max_seq_len: Optional[int] = None,
+ modules_to_fuse: Optional[dict] = None,
+ modules_to_not_convert: Optional[List] = None,
+ exllama_config: Optional[Dict[str, int]] = None,
+ **kwargs,
+ ):
+ self.quant_method = QuantizationMethod.AWQ
+
+ self.bits = bits
+ self.group_size = group_size
+ self.zero_point = zero_point
+ self.version = version
+ self.backend = backend
+ self.fuse_max_seq_len = fuse_max_seq_len
+ self.modules_to_not_convert = modules_to_not_convert
+ self.exllama_config = exllama_config
+
+ self.modules_to_fuse = modules_to_fuse
+ if do_fuse is None:
+ self.do_fuse = modules_to_fuse is not None and len(modules_to_fuse) > 0
+ else:
+ self.do_fuse = do_fuse
+ self.fuse_max_seq_len = fuse_max_seq_len
+
+ self.post_init()
+
+ def post_init(self):
+ r"""
+ Safety checker that arguments are correct
+ """
+ if self.backend not in [AwqBackendPackingMethod.AUTOAWQ, AwqBackendPackingMethod.LLMAWQ]:
+ raise ValueError(
+ f"Only supported quantization backends in {AwqBackendPackingMethod.AUTOAWQ} and {AwqBackendPackingMethod.LLMAWQ} - not recognized backend {self.backend}"
+ )
+
+ self.version = AWQLinearVersion.from_str(self.version)
+ if self.version not in [
+ AWQLinearVersion.GEMM,
+ AWQLinearVersion.GEMV,
+ AWQLinearVersion.EXLLAMA,
+ AWQLinearVersion.IPEX,
+ ]:
+ raise ValueError(
+ f"Only supported versions are in [AWQLinearVersion.GEMM, AWQLinearVersion.GEMV, AWQLinearVersion.EXLLAMA, AWQLinearVersion.IPEX] - not recognized version {self.version}"
+ )
+
+ if self.backend == AwqBackendPackingMethod.LLMAWQ:
+ # Only cuda device can run this function
+ if not (torch.cuda.is_available() or torch.xpu.is_available()):
+ raise ValueError("LLM-AWQ backend is only supported on CUDA and XPU")
+ if torch.cuda.is_available():
+ compute_capability = torch.cuda.get_device_capability()
+ major, minor = compute_capability
+ if major < 8:
+ raise ValueError("LLM-AWQ backend is only supported on CUDA GPUs with compute capability >= 8.0")
+
+ if self.do_fuse and self.fuse_max_seq_len is None:
+ raise ValueError(
+ "You cannot enable fused modules without specifying a `fuse_max_seq_len`, make sure to pass a valid `fuse_max_seq_len` for your usecase"
+ )
+
+ if self.do_fuse:
+ awq_version_supports_fusing = False
+ MIN_AWQ_VERSION = "0.1.7"
+ if is_auto_awq_available():
+ awq_version_supports_fusing = version.parse(importlib.metadata.version("autoawq")) >= version.parse(
+ MIN_AWQ_VERSION
+ )
+
+ if not awq_version_supports_fusing:
+ raise ValueError(
+ f"You current version of `autoawq` does not support module fusing, please upgrade `autoawq` package to at least {MIN_AWQ_VERSION}."
+ )
+
+ if self.modules_to_not_convert is not None:
+ awq_version_supports_non_conversion = False
+ MIN_AWQ_VERSION = "0.1.8"
+ if is_auto_awq_available():
+ awq_version_supports_non_conversion = version.parse(
+ importlib.metadata.version("autoawq")
+ ) >= version.parse(MIN_AWQ_VERSION)
+
+ if not awq_version_supports_non_conversion:
+ raise ValueError(
+ f"You current version of `autoawq` does not support module quantization skipping, please upgrade `autoawq` package to at least {MIN_AWQ_VERSION}."
+ )
+
+ if self.do_fuse and self.modules_to_fuse is not None:
+ required_keys = [
+ "hidden_size",
+ "num_attention_heads",
+ "num_key_value_heads",
+ "mlp",
+ "attention",
+ "layernorm",
+ "use_alibi",
+ ]
+ if not all(key in self.modules_to_fuse for key in required_keys):
+ raise ValueError(
+ f"Required fields are missing in the fusing mapping, required fields are {required_keys}"
+ )
+
+ if self.version == AWQLinearVersion.EXLLAMA:
+ awq_version_supports_exllama = False
+ MIN_AWQ_VERSION = "0.2.0"
+ if is_auto_awq_available():
+ awq_version_supports_exllama = version.parse(importlib.metadata.version("autoawq")) >= version.parse(
+ MIN_AWQ_VERSION
+ )
+
+ if not awq_version_supports_exllama:
+ raise ValueError(
+ f"You current version of `autoawq` does not support exllama backend, "
+ f"please upgrade `autoawq` package to at least {MIN_AWQ_VERSION}."
+ )
+
+ if self.exllama_config is None:
+ self.exllama_config = {"version": ExllamaVersion.TWO, "max_input_len": 2048, "max_batch_size": 8}
+ else:
+ if "version" not in self.exllama_config:
+ raise ValueError("`exllama_config` needs to have a `version` key.")
+ elif self.exllama_config["version"] not in [ExllamaVersion.ONE, ExllamaVersion.TWO]:
+ exllama_version = self.exllama_config["version"]
+ raise ValueError(
+ f"Only supported versions are in [ExllamaVersion.ONE, ExllamaVersion.TWO] - not recognized version {exllama_version}"
+ )
+
+ def get_loading_attributes(self):
+ attibutes_dict = copy.deepcopy(self.__dict__)
+ loading_attibutes = ["version", "do_fuse", "modules_to_fuse", "fuse_max_seq_len", "exllama_config"]
+ loading_attibutes_dict = {i: j for i, j in attibutes_dict.items() if i in loading_attibutes}
+ return loading_attibutes_dict
+
+
+@dataclass
+class AqlmConfig(QuantizationConfigMixin):
+ """
+ This is a wrapper class about `aqlm` parameters.
+
+ Args:
+ in_group_size (`int`, *optional*, defaults to 8):
+ The group size along the input dimension.
+ out_group_size (`int`, *optional*, defaults to 1):
+ The group size along the output dimension. It's recommended to always use 1.
+ num_codebooks (`int`, *optional*, defaults to 1):
+ Number of codebooks for the Additive Quantization procedure.
+ nbits_per_codebook (`int`, *optional*, defaults to 16):
+ Number of bits encoding a single codebook vector. Codebooks size is 2**nbits_per_codebook.
+ linear_weights_not_to_quantize (`Optional[List[str]]`, *optional*):
+ List of full paths of `nn.Linear` weight parameters that shall not be quantized.
+ kwargs (`Dict[str, Any]`, *optional*):
+ Additional parameters from which to initialize the configuration object.
+ """
+
+ def __init__(
+ self,
+ in_group_size: int = 8,
+ out_group_size: int = 1,
+ num_codebooks: int = 1,
+ nbits_per_codebook: int = 16,
+ linear_weights_not_to_quantize: Optional[List[str]] = None,
+ **kwargs,
+ ):
+ self.quant_method = QuantizationMethod.AQLM
+ self.in_group_size = in_group_size
+ self.out_group_size = out_group_size
+ self.num_codebooks = num_codebooks
+ self.nbits_per_codebook = nbits_per_codebook
+ self.linear_weights_not_to_quantize = linear_weights_not_to_quantize
+
+ self.post_init()
+
+ def post_init(self):
+ r"""
+ Safety checker that arguments are correct - also replaces some NoneType arguments with their default values.
+ """
+ if not isinstance(self.in_group_size, int):
+ raise TypeError("in_group_size must be a float")
+ if not isinstance(self.out_group_size, int):
+ raise TypeError("out_group_size must be a float")
+ if not isinstance(self.num_codebooks, int):
+ raise TypeError("num_codebooks must be a float")
+ if not isinstance(self.nbits_per_codebook, int):
+ raise TypeError("nbits_per_codebook must be a float")
+
+ if self.linear_weights_not_to_quantize is not None and not isinstance(
+ self.linear_weights_not_to_quantize, list
+ ):
+ raise ValueError("linear_weights_not_to_quantize must be a list of strings")
+
+ if self.linear_weights_not_to_quantize is None:
+ self.linear_weights_not_to_quantize = []
+
+
+@dataclass
+class VptqLayerConfig(QuantizationConfigMixin):
+ """
+ This is used to explain vptq config params for each layer
+ Args:
+ enable_norm (`bool`, *optional*, defaults to `True`): to control if we have scale/bias for fp-weight
+ enable_perm (`bool`, *optional*, defaults to `True`): to perm input_channel or not
+ group_num (`int`, *optional*, defaults to `1`): how many single groups for vector-quantization
+ group_size (`int`, *optional*, defaults to `-1`): depends on out-features
+ indices_as_float (`bool`, *optional*, defaults to `False`): for Finetuning
+ is_indice_packed (`bool`, *optional*, defaults to `True`): should always be True
+ num_centroids (`list`, *optional*, defaults to `[-1, -1]`): centriod numbers of clusters
+ num_res_centroids (`list`, *optional*, defaults to `[-1, -1]`): ditto for residual
+ outlier_size (`int`, *optional*, defaults to `1`): outliers
+ vector_lens (`list`, *optional*, defaults to `[-1, -1]`): centroid vector length in quantization
+ """
+
+ def __init__(
+ self,
+ enable_norm: bool = True,
+ enable_perm: bool = True,
+ group_num: int = 1,
+ group_size: int = -1,
+ in_features: int = -1,
+ indices_as_float: bool = False,
+ is_indice_packed: bool = True,
+ num_centroids: tuple = [-1, -1],
+ num_res_centroids: tuple = [-1, -1],
+ out_features: int = -1,
+ outlier_size: int = 0,
+ vector_lens: tuple = [-1, -1],
+ **kwargs,
+ ):
+ self.enable_norm = enable_norm
+ self.enable_perm = enable_perm
+ self.group_num = group_num
+ self.group_size = group_size
+ self.in_features = in_features
+ self.indices_as_float = indices_as_float
+ self.is_indice_packed = is_indice_packed
+ self.num_centroids = num_centroids
+ self.num_res_centroids = num_res_centroids
+ self.out_features = out_features
+ self.outlier_size = outlier_size
+ self.vector_lens = vector_lens
+ self.post_init()
+
+ def post_init(self):
+ r"""
+ Safety checker that arguments are correct
+ """
+ if self.is_indice_packed is False:
+ raise ValueError("is_indice_packed should always be True")
+
+
+@dataclass
+class VptqConfig(QuantizationConfigMixin):
+ """
+ This is a wrapper class about `vptq` parameters.
+
+ Args:
+ enable_proxy_error (`bool`, *optional*, defaults to `False`): calculate proxy error for each layer
+ config_for_layers (`Dict`, *optional*, defaults to `{}`): quantization params for each layer
+ shared_layer_config (`Dict`, *optional*, defaults to `{}`): shared quantization params among layers
+ modules_to_not_convert (`list`, *optional*, default to `None`):
+ The list of modules to not quantize, useful for quantizing models that explicitly require to have
+ some modules left in their original precision (e.g. Whisper encoder, Llava encoder, Mixtral gate layers).
+ kwargs (`Dict[str, Any]`, *optional*):
+ Additional parameters from which to initialize the configuration object.
+ """
+
+ def __init__(
+ self,
+ enable_proxy_error: bool = False,
+ config_for_layers: Dict[str, Any] = {},
+ shared_layer_config: Dict[str, Any] = {},
+ modules_to_not_convert: Optional[List] = None,
+ **kwargs,
+ ):
+ self.quant_method = QuantizationMethod.VPTQ
+ self.enable_proxy_error = enable_proxy_error
+ self.config_for_layers: Dict[str, Any] = config_for_layers
+ self.shared_layer_config: Dict[str, Any] = shared_layer_config
+ self.modules_to_not_convert = modules_to_not_convert
+ self.post_init()
+
+ def post_init(self):
+ r"""
+ Safety checker that arguments are correct
+ """
+ for layer_name, layer_param in self.config_for_layers.items():
+ VptqLayerConfig(**layer_param)
+ if self.enable_proxy_error is True:
+ raise ValueError("enable_proxy_error should always be False until we support training")
+
+
+@dataclass
+class QuantoConfig(QuantizationConfigMixin):
+ """
+ This is a wrapper class about all possible attributes and features that you can play with a model that has been
+ loaded using `quanto`.
+
+ Args:
+ weights (`str`, *optional*, defaults to `"int8"`):
+ The target dtype for the weights after quantization. Supported values are ("float8","int8","int4","int2")
+ activations (`str`, *optional*):
+ The target dtype for the activations after quantization. Supported values are (None,"int8","float8")
+ modules_to_not_convert (`list`, *optional*, default to `None`):
+ The list of modules to not quantize, useful for quantizing models that explicitly require to have
+ some modules left in their original precision (e.g. Whisper encoder, Llava encoder, Mixtral gate layers).
+ """
+
+ def __init__(
+ self,
+ weights="int8",
+ activations=None,
+ modules_to_not_convert: Optional[List] = None,
+ **kwargs,
+ ):
+ self.quant_method = QuantizationMethod.QUANTO
+ self.weights = weights
+ self.activations = activations
+ self.modules_to_not_convert = modules_to_not_convert
+ self.post_init()
+
+ def post_init(self):
+ r"""
+ Safety checker that arguments are correct
+ """
+ accepted_weights = ["float8", "int8", "int4", "int2"]
+ accepted_activations = [None, "int8", "float8"]
+ if self.weights not in accepted_weights:
+ raise ValueError(f"Only support weights in {accepted_weights} but found {self.weights}")
+ if self.activations not in accepted_activations:
+ raise ValueError(f"Only support weights in {accepted_activations} but found {self.activations}")
+
+
+@dataclass
+class EetqConfig(QuantizationConfigMixin):
+ """
+ This is a wrapper class about all possible attributes and features that you can play with a model that has been
+ loaded using `eetq`.
+
+ Args:
+ weights (`str`, *optional*, defaults to `"int8"`):
+ The target dtype for the weights. Supported value is only "int8"
+ modules_to_not_convert (`list`, *optional*, default to `None`):
+ The list of modules to not quantize, useful for quantizing models that explicitly require to have
+ some modules left in their original precision.
+ """
+
+ def __init__(
+ self,
+ weights: str = "int8",
+ modules_to_not_convert: Optional[List] = None,
+ **kwargs,
+ ):
+ self.quant_method = QuantizationMethod.EETQ
+ self.weights = weights
+ self.modules_to_not_convert = modules_to_not_convert
+ self.post_init()
+
+ def post_init(self):
+ r"""
+ Safety checker that arguments are correct
+ """
+ accepted_weights = ["int8"]
+ if self.weights not in accepted_weights:
+ raise ValueError(f"Only support weights in {accepted_weights} but found {self.weights}")
+
+
+class CompressedTensorsConfig(QuantizationConfigMixin):
+ """
+ This is a wrapper class that handles compressed-tensors quantization config options.
+ It is a wrapper around `compressed_tensors.QuantizationConfig`
+ Args:
+ config_groups (`typing.Dict[str, typing.Union[ForwardRef('QuantizationScheme'), typing.List[str]]]`, *optional*):
+ dictionary mapping group name to a quantization scheme definition
+ format (`str`, *optional*, defaults to `"dense"`):
+ format the model is represented as. Set `run_compressed` True to execute model as the
+ compressed format if not `dense`
+ quantization_status (`QuantizationStatus`, *optional*, defaults to `"initialized"`):
+ status of model in the quantization lifecycle, ie 'initialized', 'calibration', 'frozen'
+ kv_cache_scheme (`typing.Union[QuantizationArgs, NoneType]`, *optional*):
+ specifies quantization of the kv cache. If None, kv cache is not quantized.
+ global_compression_ratio (`typing.Union[float, NoneType]`, *optional*):
+ 0-1 float percentage of model compression
+ ignore (`typing.Union[typing.List[str], NoneType]`, *optional*):
+ layer names or types to not quantize, supports regex prefixed by 're:'
+ sparsity_config (`typing.Dict[str, typing.Any]`, *optional*):
+ configuration for sparsity compression
+ quant_method (`str`, *optional*, defaults to `"compressed-tensors"`):
+ do not override, should be compressed-tensors
+ run_compressed (`bool`, *optional*, defaults to `True`): alter submodules (usually linear) in order to
+ emulate compressed model execution if True, otherwise use default submodule
+ """
+
+ def __init__(
+ self,
+ config_groups: Dict[str, Union["QuantizationScheme", List[str]]] = None, # noqa: F821
+ format: str = "dense",
+ quantization_status: "QuantizationStatus" = "initialized", # noqa: F821
+ kv_cache_scheme: Optional["QuantizationArgs"] = None, # noqa: F821
+ global_compression_ratio: Optional[float] = None,
+ ignore: Optional[List[str]] = None,
+ sparsity_config: Dict[str, Any] = None,
+ quant_method: str = "compressed-tensors",
+ run_compressed: bool = True,
+ **kwargs,
+ ):
+ if is_compressed_tensors_available():
+ from compressed_tensors.config import SparsityCompressionConfig
+ from compressed_tensors.quantization import QuantizationConfig
+ else:
+ raise ImportError(
+ "compressed_tensors is not installed and is required for compressed-tensors quantization. Please install it with `pip install compressed-tensors`."
+ )
+ self.quantization_config = None
+ self.sparsity_config = None
+
+ self.run_compressed = run_compressed
+
+ # parse from dict to load nested QuantizationScheme objects
+ if config_groups or kv_cache_scheme:
+ self.quantization_config = QuantizationConfig.model_validate(
+ {
+ "config_groups": config_groups,
+ "quant_method": quant_method,
+ "format": format,
+ "quantization_status": quantization_status,
+ "kv_cache_scheme": kv_cache_scheme,
+ "global_compression_ratio": global_compression_ratio,
+ "ignore": ignore,
+ "run_compressed": run_compressed,
+ **kwargs,
+ }
+ )
+
+ if sparsity_config:
+ self.sparsity_config = SparsityCompressionConfig.load_from_registry(
+ sparsity_config.get("format"), **sparsity_config
+ )
+
+ self.quant_method = QuantizationMethod.COMPRESSED_TENSORS
+
+ def post_init(self):
+ if self.run_compressed:
+ if self.is_sparsification_compressed:
+ logger.warn(
+ "`run_compressed` is only supported for quantized_compressed models"
+ " and not for sparsified models. Setting `run_compressed=False`"
+ )
+ self.run_compressed = False
+ elif not self.is_quantization_compressed:
+ logger.warn("`run_compressed` is only supported for compressed models. Setting `run_compressed=False`")
+ self.run_compressed = False
+
+ @classmethod
+ def from_dict(cls, config_dict, return_unused_kwargs=False, **kwargs):
+ """
+ Instantiates a [`CompressedTensorsConfig`] from a Python dictionary of parameters.
+ Optionally unwraps any args from the nested quantization_config
+
+ Args:
+ config_dict (`Dict[str, Any]`):
+ Dictionary that will be used to instantiate the configuration object.
+ return_unused_kwargs (`bool`,*optional*, defaults to `False`):
+ Whether or not to return a list of unused keyword arguments. Used for `from_pretrained` method in
+ `PreTrainedModel`.
+ kwargs (`Dict[str, Any]`):
+ Additional parameters from which to initialize the configuration object.
+
+ Returns:
+ [`QuantizationConfigMixin`]: The configuration object instantiated from those parameters.
+
+ """
+
+ if "quantization_config" in config_dict:
+ config_dict = dict(
+ sparsity_config=config_dict.get("sparsity_config"),
+ **config_dict["quantization_config"],
+ )
+
+ return super().from_dict(config_dict, return_unused_kwargs=return_unused_kwargs, **kwargs)
+
+ def to_dict(self) -> Dict[str, Any]:
+ """
+ Quantization config to be added to config.json
+
+ Serializes this instance to a Python dictionary. Returns:
+ `Dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance.
+ """
+ quantization_config = {}
+ if self.quantization_config is not None:
+ quantization_config = self.quantization_config.dict()
+ else:
+ quantization_config["quant_method"] = QuantizationMethod.COMPRESSED_TENSORS
+
+ if self.sparsity_config is not None:
+ quantization_config["sparsity_config"] = self.sparsity_config.dict()
+ else:
+ quantization_config["sparsity_config"] = {}
+
+ return quantization_config
+
+ def to_diff_dict(self) -> Dict[str, Any]:
+ """
+ Removes all attributes from config which correspond to the default config attributes for better readability and
+ serializes to a Python dictionary.
+ Returns:
+ `Dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance,
+ """
+ config_dict = self.to_dict()
+
+ # get the default config dict
+ default_config_dict = CompressedTensorsConfig().to_dict()
+
+ serializable_config_dict = {}
+
+ # only serialize values that differ from the default config
+ for key, value in config_dict.items():
+ if key not in default_config_dict or value != default_config_dict[key]:
+ serializable_config_dict[key] = value
+
+ return serializable_config_dict
+
+ def get_loading_attributes(self):
+ return {"run_compressed": self.run_compressed}
+
+ @property
+ def is_quantized(self):
+ return bool(self.quantization_config) and bool(self.quantization_config.config_groups)
+
+ @property
+ def is_quantization_compressed(self):
+ from compressed_tensors.quantization import QuantizationStatus
+
+ return self.is_quantized and self.quantization_config.quantization_status == QuantizationStatus.COMPRESSED
+
+ @property
+ def is_sparsification_compressed(self):
+ from compressed_tensors.config import (
+ CompressionFormat,
+ SparsityCompressionConfig,
+ )
+
+ return (
+ isinstance(self.sparsity_config, SparsityCompressionConfig)
+ and self.sparsity_config.format != CompressionFormat.dense.value
+ )
+
+
+@dataclass
+class FbgemmFp8Config(QuantizationConfigMixin):
+ """
+ This is a wrapper class about all possible attributes and features that you can play with a model that has been
+ loaded using fbgemm fp8 quantization.
+
+ Args:
+ activation_scale_ub (`float`, *optional*, defaults to 1200.0):
+ The activation scale upper bound. This is used when quantizing the input activation.
+ modules_to_not_convert (`list`, *optional*, default to `None`):
+ The list of modules to not quantize, useful for quantizing models that explicitly require to have
+ some modules left in their original precision.
+ """
+
+ def __init__(
+ self,
+ activation_scale_ub: float = 1200.0,
+ modules_to_not_convert: Optional[List] = None,
+ **kwargs,
+ ):
+ self.quant_method = QuantizationMethod.FBGEMM_FP8
+ self.activation_scale_ub = activation_scale_ub
+ self.modules_to_not_convert = modules_to_not_convert
+
+ def get_loading_attributes(self):
+ attibutes_dict = copy.deepcopy(self.__dict__)
+ loading_attibutes = ["activation_scale_ub"]
+ loading_attibutes_dict = {i: j for i, j in attibutes_dict.items() if i in loading_attibutes}
+ return loading_attibutes_dict
+
+
+@dataclass
+class HiggsConfig(QuantizationConfigMixin):
+ """
+ HiggsConfig is a configuration class for quantization using the HIGGS method.
+
+ Args:
+ bits (int, *optional*, defaults to 4):
+ Number of bits to use for quantization. Can be 2, 3 or 4. Default is 4.
+ p (int, *optional*, defaults to 2):
+ Quantization grid dimension. 1 and 2 are supported. 2 is always better in practice. Default is 2.
+ modules_to_not_convert (`list`, *optional*, default to ["lm_head"]):
+ List of linear layers that should not be quantized.
+ hadamard_size (int, *optional*, defaults to 512):
+ Hadamard size for the HIGGS method. Default is 512. Input dimension of matrices is padded to this value. Decreasing this below 512 will reduce the quality of the quantization.
+ group_size (int, *optional*, defaults to 256):
+ Group size for the HIGGS method. Can be 64, 128 or 256. Decreasing it barely affects the performance. Default is 256. Must be a divisor of hadamard_size.
+ tune_metadata ('dict', *optional*, defaults to {}):
+ Module-wise metadata (gemm block shapes, GPU metadata, etc.) for saving the kernel tuning results. Default is an empty dictionary. Is set automatically during tuning.
+ """
+
+ def __init__(
+ self,
+ bits: int = 4,
+ p: int = 2,
+ modules_to_not_convert: Optional[List[str]] = None,
+ hadamard_size: int = 512,
+ group_size: int = 256,
+ tune_metadata: Optional[Dict[str, Any]] = None,
+ **kwargs,
+ ):
+ if tune_metadata is None:
+ tune_metadata = {}
+ self.quant_method = QuantizationMethod.HIGGS
+ self.bits = bits
+ self.p = p
+ self.modules_to_not_convert = modules_to_not_convert
+ self.hadamard_size = hadamard_size
+ self.group_size = group_size
+ self.tune_metadata = tune_metadata
+
+ self.post_init()
+
+ def post_init(self):
+ r"""
+ Safety checker that arguments are correct - also replaces some NoneType arguments with their default values.
+ """
+ if self.bits not in [2, 3, 4]:
+ raise ValueError("bits must be 2, 3, or 4")
+ if self.p not in [1, 2]:
+ raise ValueError("p must be 1 or 2. 2 is always better in practice")
+ if self.group_size not in [64, 128, 256]:
+ raise ValueError("group_size must be 64, 128, or 256")
+ if self.hadamard_size % self.group_size != 0:
+ raise ValueError("hadamard_size must be divisible by group_size")
+
+
+@dataclass
+class TorchAoConfig(QuantizationConfigMixin):
+ quant_method: QuantizationMethod
+ quant_type: Union[str, "AOBaseConfig"] # noqa: F821
+ modules_to_not_convert: Optional[List]
+ quant_type_kwargs: Dict[str, Any]
+
+ """This is a config class for torchao quantization/sparsity techniques.
+
+ Args:
+ quant_type (`Union[str, AOBaseConfig]`):
+ The type of quantization we want to use. Can be either:
+ - A string: currently supporting: `int4_weight_only`, `int8_weight_only` and `int8_dynamic_activation_int8_weight`.
+ - An AOBaseConfig instance: for more advanced configuration options.
+ modules_to_not_convert (`list`, *optional*, default to `None`):
+ The list of modules to not quantize, useful for quantizing models that explicitly require to have
+ some modules left in their original precision.
+ kwargs (`Dict[str, Any]`, *optional*):
+ The keyword arguments for the chosen type of quantization, for example, int4_weight_only quantization supports two keyword arguments
+ `group_size` and `inner_k_tiles` currently. More API examples and documentation of arguments can be found in
+ https://github.com/pytorch/ao/tree/main/torchao/quantization#other-available-quantization-techniques
+
+ Example:
+
+ ```python
+ # AOBaseConfig-based configuration
+ config = Int4WeightOnlyConfig(group_size=32)
+ quantization_config = TorchAoConfig(config)
+ model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cuda", torch_dtype=torch.bfloat16, quantization_config=quantization_config)
+
+ # String-based configuration
+ quantization_config = TorchAoConfig("int4_weight_only", group_size=32)
+ # int4_weight_only quant is only working with *torch.bfloat16* dtype right now
+ model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cuda", torch_dtype=torch.bfloat16, quantization_config=quantization_config)
+
+ # autoquant
+ # `autoquant` is a convenient way for users to search for the best quantization for each layer
+ # `min_sqnr` is an option to control the accuracy of the model, higher value means the model is more
+ # accurate, we can start with 30 and adjust it to larger or smaller (e.g. 40, 20)
+ # defaults to None, which means we'll try to get the best performing quantized model without
+ # considering accuracy
+ quantization_config = TorchAoConfig("autoquant", min_sqnr=30)
+ model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cuda", torch_dtype=torch.bfloat16, quantization_config=quantization_config)
+ # run through example inputs, quantization methods will be selected based on the shape of example input
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
+ input_text = "What are we having for dinner?"
+ input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")
+ MAX_NEW_TOKENS = 1000
+ model.generate(**input_ids, max_new_tokens=MAX_NEW_TOKENS, cache_implementation="static")
+ # manually ran finalize_autoquant if needed
+ if hasattr(quantized_model, "finalize_autoquant"):
+ print("finalizing autoquant")
+ quantized_model.finalize_autoquant()
+
+ ```
+ """
+
+ def __init__(
+ self,
+ quant_type: Union[str, "AOBaseConfig"], # noqa: F821
+ modules_to_not_convert: Optional[List] = None,
+ **kwargs,
+ ):
+ self.quant_method = QuantizationMethod.TORCHAO
+ self.quant_type = quant_type
+ self.modules_to_not_convert = modules_to_not_convert
+ self.quant_type_kwargs = kwargs.get("quant_type_kwargs", kwargs)
+ self.post_init()
+
+ @staticmethod
+ def _get_ao_version() -> version.Version:
+ """Centralized check for TorchAO availability and version requirements."""
+ if not is_torchao_available():
+ raise ValueError("TorchAoConfig requires torchao to be installed. Install with `pip install torchao`")
+
+ return version.parse(importlib.metadata.version("torchao"))
+
+ def post_init(self):
+ """Validate configuration and set defaults."""
+ ao_version = self._get_ao_version()
+
+ # Handle quant_type based on type and version
+ if isinstance(self.quant_type, str):
+ self._validate_string_quant_type()
+ elif ao_version > version.parse("0.9.0"):
+ from torchao.quantization.quant_api import AOBaseConfig
+
+ if not isinstance(self.quant_type, AOBaseConfig):
+ raise ValueError(
+ f"quant_type must be either a string or an AOBaseConfig instance, got {type(self.quant_type)}"
+ )
+ else:
+ raise ValueError(
+ f"In torchao <= 0.9.0, quant_type must be a string. Got {type(self.quant_type)}. "
+ f"Please upgrade to torchao > 0.9.0 to use AOBaseConfig instances."
+ )
+
+ def _validate_string_quant_type(self):
+ """Validate string quant_type and its kwargs."""
+ methods = self._get_torchao_quant_type_to_method()
+
+ if self.quant_type not in methods:
+ raise ValueError(
+ f"Unsupported string quantization type: {self.quant_type}. "
+ f"Supported types: {', '.join(methods.keys())}"
+ )
+
+ # Validate kwargs against method signature
+ method = methods[self.quant_type]
+ sig = signature(method)
+ valid_kwargs = {
+ param.name
+ for param in sig.parameters.values()
+ if param.kind in [Parameter.KEYWORD_ONLY, Parameter.POSITIONAL_OR_KEYWORD]
+ }
+
+ invalid_kwargs = set(self.quant_type_kwargs) - valid_kwargs
+ if invalid_kwargs:
+ raise ValueError(
+ f"Unexpected keyword arg for {self.quant_type}: {', '.join(invalid_kwargs)}. "
+ f"Valid kwargs: {', '.join(valid_kwargs)}"
+ )
+
+ def _get_torchao_quant_type_to_method(self):
+ """Get mapping of quant_type strings to their corresponding methods."""
+ from torchao.quantization import (
+ autoquant,
+ int4_weight_only,
+ int8_dynamic_activation_int8_weight,
+ int8_weight_only,
+ )
+
+ return {
+ "int4_weight_only": int4_weight_only,
+ "int8_weight_only": int8_weight_only,
+ "int8_dynamic_activation_int8_weight": int8_dynamic_activation_int8_weight,
+ "autoquant": autoquant,
+ }
+
+ def get_apply_tensor_subclass(self):
+ """Create the appropriate quantization method based on configuration."""
+ if isinstance(self.quant_type, str):
+ methods = self._get_torchao_quant_type_to_method()
+ quant_type_kwargs = self.quant_type_kwargs.copy()
+ if (
+ not torch.cuda.is_available()
+ and is_torchao_available()
+ and self.quant_type == "int4_weight_only"
+ and version.parse(importlib.metadata.version("torchao")) >= version.parse("0.8.0")
+ ):
+ from torchao.dtypes import Int4CPULayout
+
+ quant_type_kwargs["layout"] = Int4CPULayout()
+
+ return methods[self.quant_type](**quant_type_kwargs)
+ else:
+ return self.quant_type
+
+ def to_dict(self):
+ """Convert configuration to a dictionary."""
+ d = super().to_dict()
+
+ if isinstance(self.quant_type, str):
+ # Handle layout serialization if present
+ if "quant_type_kwargs" in d and "layout" in d["quant_type_kwargs"]:
+ d["quant_type_kwargs"]["layout"] = dataclasses.asdict(d["quant_type_kwargs"]["layout"])
+ else:
+ # Handle AOBaseConfig serialization
+ from torchao.core.config import config_to_dict
+
+ # For now we assume there is 1 config per Transfomer, however in the future
+ # We may want to support a config per fqn.
+ d["quant_type"] = {"default": config_to_dict(self.quant_type)}
+
+ return d
+
+ @classmethod
+ def from_dict(cls, config_dict, return_unused_kwargs=False, **kwargs):
+ """Create configuration from a dictionary."""
+ ao_verison = cls._get_ao_version()
+ assert ao_verison > version.parse("0.9.0"), "TorchAoConfig requires torchao > 0.9.0 for construction from dict"
+ config_dict = config_dict.copy()
+ quant_type = config_dict.pop("quant_type")
+ # Check if we only have one key which is "default"
+ # In the future we may update this
+ assert len(quant_type) == 1 and "default" in quant_type, (
+ "Expected only one key 'default' in quant_type dictionary"
+ )
+ quant_type = quant_type["default"]
+
+ # Deserialize quant_type if needed
+ from torchao.core.config import config_from_dict
+
+ quant_type = config_from_dict(quant_type)
+
+ return cls(quant_type=quant_type, **config_dict)
+
+
+@dataclass
+class BitNetConfig(QuantizationConfigMixin):
+ def __init__(
+ self,
+ modules_to_not_convert: Optional[List] = None,
+ **kwargs,
+ ):
+ self.quant_method = QuantizationMethod.BITNET
+ self.modules_to_not_convert = modules_to_not_convert
+ self.post_init()
+
+ def post_init(self):
+ r"""
+ Safety checker that arguments are correct
+ """
+ pass
+
+
+@dataclass
+class SpQRConfig(QuantizationConfigMixin):
+ """
+ This is a wrapper class about `spqr` parameters. Refer to the original publication for more details.
+
+ Args:
+ bits (`int`, *optional*, defaults to 3):
+ Specifies the bit count for the weights and first order zero-points and scales.
+ Currently only bits = 3 is supported.
+ beta1 (`int`, *optional*, defaults to 16):
+ SpQR tile width. Currently only beta1 = 16 is supported.
+ beta2 (`int`, *optional*, defaults to 16):
+ SpQR tile height. Currently only beta2 = 16 is supported.
+ shapes (`Optional`, *optional*):
+ A dictionary holding the shape of each object. We need this because it's impossible
+ to deduce the exact size of the parameters just from bits, beta1, beta2.
+ modules_to_not_convert (`Optional[List[str]]`, *optional*):
+ Optionally, provides a list of full paths of `nn.Linear` weight parameters that shall not be quantized.
+ Defaults to None.
+ kwargs (`Dict[str, Any]`, *optional*):
+ Additional parameters from which to initialize the configuration object.
+ """
+
+ def __init__(
+ self,
+ bits: int = 3,
+ beta1: int = 16,
+ beta2: int = 16,
+ shapes: Optional[Dict[str, int]] = None,
+ modules_to_not_convert: Optional[List[str]] = None,
+ **kwargs,
+ ):
+ if shapes is None:
+ shapes = {}
+ self.shapes = shapes
+ self.quant_method = QuantizationMethod.SPQR
+ self.bits = bits
+ self.beta1 = beta1
+ self.beta2 = beta2
+ self.modules_to_not_convert = modules_to_not_convert
+ self.post_init()
+
+ def post_init(self):
+ r"""
+ Safety checker that arguments are correct - also replaces some NoneType arguments with their default values.
+ """
+ if not isinstance(self.bits, int):
+ raise TypeError("bits must be an int")
+ if not isinstance(self.beta1, int):
+ raise TypeError("beta1 must be an int")
+ if not isinstance(self.beta2, int):
+ raise TypeError("beta2 must be an int")
+
+ if self.bits != 3:
+ raise ValueError("SpQR currently only supports bits = 3")
+ if self.beta1 != 16:
+ raise ValueError("SpQR currently only supports beta1 = 16")
+ if self.beta2 != 16:
+ raise ValueError("SpQR currently only supports beta2 = 16")
+ if not isinstance(self.shapes, dict):
+ raise TypeError("shapes must be a dict")
+
+
+@dataclass
+class FineGrainedFP8Config(QuantizationConfigMixin):
+ """
+ FineGrainedFP8Config is a configuration class for fine-grained FP8 quantization used mainly for deepseek models.
+
+ Args:
+ activation_scheme (`str`, *optional*, defaults to `"dynamic"`):
+ The scheme used for activation, the defaults and only support scheme for now is "dynamic".
+ weight_block_size (`typing.Tuple[int, int]`, *optional*, defaults to `(128, 128)`):
+ The size of the weight blocks for quantization, default is (128, 128).
+ modules_to_not_convert (`list`, *optional*):
+ A list of module names that should not be converted during quantization.
+ """
+
+ def __init__(
+ self,
+ activation_scheme: str = "dynamic",
+ weight_block_size: Tuple[int, int] = (128, 128),
+ modules_to_not_convert: Optional[List] = None,
+ **kwargs,
+ ):
+ self.quant_method = QuantizationMethod.FP8
+ self.modules_to_not_convert = modules_to_not_convert
+ self.activation_scheme = activation_scheme
+ self.weight_block_size = weight_block_size
+ self.post_init()
+
+ def post_init(self):
+ r"""
+ Safety checker that arguments are correct
+ """
+ self.activation_scheme = self.activation_scheme.lower()
+ if self.activation_scheme not in ["dynamic"]:
+ raise ValueError(f"Activation scheme {self.activation_scheme} not supported")
+ if len(self.weight_block_size) != 2:
+ raise ValueError("weight_block_size must be a tuple of two integers")
+ if self.weight_block_size[0] <= 0 or self.weight_block_size[1] <= 0:
+ raise ValueError("weight_block_size must be a tuple of two positive integers")
+
+
+class QuarkConfig(QuantizationConfigMixin):
+ def __init__(
+ self,
+ **kwargs,
+ ):
+ if is_torch_available() and is_quark_available():
+ from quark import __version__ as quark_version
+ from quark.torch.export.config.config import JsonExporterConfig
+ from quark.torch.export.main_export.quant_config_parser import QuantConfigParser
+ from quark.torch.quantization.config.config import Config
+
+ # This might be e.g. `"fp8"` or `"awq"`.
+ self.custom_mode = kwargs["quant_method"]
+ self.legacy = "export" not in kwargs
+
+ if self.custom_mode in ["awq", "fp8"]:
+ # Legacy (quark<1.0) or custom export.
+ self.quant_config = QuantConfigParser.from_custom_config(kwargs, is_bias_quantized=False)
+ self.json_export_config = JsonExporterConfig()
+ else:
+ self.quant_config = Config.from_dict(kwargs)
+
+ if "export" in kwargs:
+ # TODO: Remove this check once configuration version is handled natively by Quark.
+ if "min_kv_scale" in kwargs["export"] and version.parse(quark_version) < version.parse("0.8"):
+ min_kv_scale = kwargs["export"].pop("min_kv_scale")
+ logger.warning(
+ f"The parameter `min_kv_scale={min_kv_scale}` was found in the model config.json's `quantization_config.export` configuration, but this parameter is supported only for quark>=0.8. Ignoring this configuration parameter. Please update the `amd-quark` package."
+ )
+
+ self.json_export_config = JsonExporterConfig(**kwargs["export"])
+ else:
+ # Legacy (quark<1.0) or custom export.
+ self.json_export_config = JsonExporterConfig()
+
+ self.quant_method = QuantizationMethod.QUARK
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_aqlm.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_aqlm.py
new file mode 100644
index 0000000000000000000000000000000000000000..9d1d6f7e89f1e9aa2454a1fd56c81524660f9079
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_aqlm.py
@@ -0,0 +1,97 @@
+# Copyright 2024 The HuggingFace Inc. 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.
+import importlib
+from typing import TYPE_CHECKING, Optional
+
+from packaging import version
+
+from .base import HfQuantizer
+
+
+if TYPE_CHECKING:
+ from ..modeling_utils import PreTrainedModel
+
+from ..integrations import replace_with_aqlm_linear
+from ..utils import is_accelerate_available, is_aqlm_available, is_torch_available, logging
+from ..utils.quantization_config import QuantizationConfigMixin
+
+
+if is_torch_available():
+ import torch
+
+logger = logging.get_logger(__name__)
+
+
+class AqlmHfQuantizer(HfQuantizer):
+ """
+ Quantizer of the AQLM method. Enables the loading of prequantized models.
+ """
+
+ requires_calibration = True
+ required_packages = ["aqlm"]
+ optimum_quantizer = None
+
+ def __init__(self, quantization_config: QuantizationConfigMixin, **kwargs):
+ super().__init__(quantization_config, **kwargs)
+ self.quantization_config = quantization_config
+
+ def validate_environment(self, *args, **kwargs):
+ if not is_accelerate_available():
+ raise ImportError("Using `aqlm` quantization requires Accelerate: `pip install accelerate`")
+
+ if not is_aqlm_available():
+ raise ImportError("Using `aqlm` quantization requires AQLM: `pip install aqlm[gpu,cpu]`")
+
+ def update_torch_dtype(self, torch_dtype: "torch.dtype") -> "torch.dtype":
+ if torch_dtype is None:
+ if torch.cuda.is_available():
+ torch_dtype = torch.float16
+ logger.info(
+ "CUDA available. Assuming AQLM inference on GPU and loading the model in `torch.float16`. To overwrite it, set `torch_dtype` manually."
+ )
+ else:
+ torch_dtype = torch.float32
+ logger.info(
+ "CUDA is unavailable. Assuming AQLM inference on CPU and loading the model in `torch.float32`. To overwrite it, set `torch_dtype` manually."
+ )
+ return torch_dtype
+
+ def _process_model_before_weight_loading(
+ self,
+ model: "PreTrainedModel",
+ **kwargs,
+ ):
+ replace_with_aqlm_linear(
+ model,
+ quantization_config=self.quantization_config,
+ linear_weights_not_to_quantize=self.quantization_config.linear_weights_not_to_quantize,
+ )
+ model.config.quantization_config = self.quantization_config
+
+ def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs):
+ return model
+
+ @property
+ def is_trainable(self, model: Optional["PreTrainedModel"] = None):
+ aqlm_supports_training = version.parse(importlib.metadata.version("aqlm")) >= version.parse("1.0.2")
+ if aqlm_supports_training:
+ return True
+ else:
+ logger.warning(
+ f"Currently installed `aqlm` version ({importlib.metadata.version('aqlm')}) doesn't support training. If you wish to train a quantized model, please update `aqlm` with `pip install aqlm>=1.0.2`"
+ )
+ return False
+
+ def is_serializable(self, safe_serialization=None):
+ return True
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_awq.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_awq.py
new file mode 100644
index 0000000000000000000000000000000000000000..28460ac38edb5e899689df7259d95b5d733bce6d
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_awq.py
@@ -0,0 +1,153 @@
+# Copyright 2024 The HuggingFace Inc. 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.
+import importlib.metadata
+from typing import TYPE_CHECKING, List, Optional
+
+from packaging import version
+
+from .base import HfQuantizer
+
+
+if TYPE_CHECKING:
+ from ..modeling_utils import PreTrainedModel
+
+from ..utils import is_accelerate_available, is_auto_awq_available, is_torch_available, logging
+from ..utils.quantization_config import AWQLinearVersion
+
+
+if is_torch_available():
+ import torch
+
+logger = logging.get_logger(__name__)
+
+
+class AwqQuantizer(HfQuantizer):
+ """
+ 4-bit quantization for Activation-aware Weight Quantization(AWQ) (https://arxiv.org/abs/2306.00978)
+ """
+
+ # AWQ requires data callibration - we support only inference
+ requires_calibration = True
+
+ required_packages = ["awq", "accelerate"]
+
+ def __init__(self, quantization_config, **kwargs):
+ super().__init__(quantization_config, **kwargs)
+
+ def validate_environment(self, device_map, **kwargs):
+ if not is_auto_awq_available():
+ raise ImportError("Loading an AWQ quantized model requires auto-awq library (`pip install autoawq`)")
+
+ if not is_accelerate_available():
+ raise ImportError("Loading an AWQ quantized model requires accelerate (`pip install accelerate`)")
+
+ if self.quantization_config.version == AWQLinearVersion.GEMM and not torch.cuda.is_available():
+ logger.warning_once("No CUDA found, replace GEMM with IPEX version to support non-cuda AWQ model.")
+ self.quantization_config.version = AWQLinearVersion.IPEX
+
+ if self.quantization_config.version == AWQLinearVersion.IPEX:
+ if version.parse(importlib.metadata.version("autoawq")) < version.parse("0.2.6"):
+ raise RuntimeError(
+ "To use IPEX backend, you need autoawq>0.2.6. Please install the latest version or from source."
+ )
+ if device_map is None:
+ logger.warning_once(
+ "You have loaded an AWQ model without setting device_map, please set 'cpu' or 'xpu' or 'auto'"
+ )
+ elif isinstance(device_map, dict) and "disk" in device_map.values():
+ raise ValueError(
+ "You are attempting to load an IPEX version AWQ model with a device_map that contains disk device."
+ " This is not supported. Please make sure only cpu and xpu in the device_map."
+ )
+ else:
+ if not torch.cuda.is_available():
+ raise RuntimeError(
+ "GPU is required to run AWQ quantized model. You can use IPEX version AWQ if you have an Intel CPU"
+ )
+
+ if device_map is None:
+ logger.warning_once(
+ "You have loaded an AWQ model on CPU and have a CUDA device available, make sure to set "
+ "your model on a GPU device in order to run your model."
+ )
+ elif device_map is not None:
+ if isinstance(device_map, dict) and ("cpu" in device_map.values() or "disk" in device_map.values()):
+ raise ValueError(
+ "You are attempting to load an AWQ model with a device_map that contains a CPU or disk device."
+ " This is not supported. Please remove the CPU or disk device from the device_map."
+ )
+
+ def update_torch_dtype(self, torch_dtype):
+ if torch_dtype is None:
+ torch_dtype = torch.float16
+ logger.info("Loading the model in `torch.float16`. To overwrite it, set `torch_dtype` manually.")
+ elif torch_dtype != torch.float16:
+ logger.warning("We suggest you to set `torch_dtype=torch.float16` for better efficiency with AWQ.")
+ return torch_dtype
+
+ def _process_model_before_weight_loading(
+ self, model: "PreTrainedModel", keep_in_fp32_modules: Optional[List[str]] = None, **kwargs
+ ):
+ from ..integrations import replace_quantization_scales, replace_with_awq_linear
+
+ self.modules_to_not_convert = self.get_modules_to_not_convert(
+ model, self.quantization_config.modules_to_not_convert, keep_in_fp32_modules
+ )
+
+ model, has_been_replaced = replace_with_awq_linear(
+ model, quantization_config=self.quantization_config, modules_to_not_convert=self.modules_to_not_convert
+ )
+
+ model = replace_quantization_scales(model, model.config.model_type)
+
+ if not has_been_replaced:
+ logger.warning(
+ "You are loading an AWQ model but no linear modules were found in your model."
+ " Please double check your model architecture, or submit an issue on github if you think this is a bug."
+ )
+
+ def _process_model_after_weight_loading(self, model, **kwargs):
+ if self.quantization_config.do_fuse:
+ from ..integrations import fuse_awq_modules
+
+ model = fuse_awq_modules(model, self.quantization_config)
+ model._awq_is_fused = True # TODO: consider storing this flag in model.config instead
+
+ if self.quantization_config.version == AWQLinearVersion.EXLLAMA:
+ from ..integrations import post_init_awq_exllama_modules
+
+ model = post_init_awq_exllama_modules(model, self.quantization_config.exllama_config)
+
+ if self.quantization_config.version == AWQLinearVersion.IPEX:
+ from ..integrations import post_init_awq_ipex_modules
+
+ model = post_init_awq_ipex_modules(model)
+
+ def is_serializable(self, safe_serialization=None):
+ # AWQ through auto-awq has been always serializable, except if the model is fused.
+ if self.quantization_config.do_fuse:
+ logger.warning("You cannot save an AWQ model that uses fused modules!")
+ return False
+
+ if self.quantization_config.version == AWQLinearVersion.EXLLAMA:
+ logger.warning("You cannot save an AWQ model that uses Exllama backend!")
+ return False
+
+ return True
+
+ @property
+ def is_trainable(self):
+ # AWQ supports PEFT fine-tuning from version 0.2.0
+ MIN_AWQ_VERSION_FOR_PEFT = "0.2.0"
+ return version.parse(importlib.metadata.version("autoawq")) >= version.parse(MIN_AWQ_VERSION_FOR_PEFT)
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_bitnet.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_bitnet.py
new file mode 100644
index 0000000000000000000000000000000000000000..e56bb161ac449853f0fb12bc86aed1a9c75b88ee
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_bitnet.py
@@ -0,0 +1,113 @@
+# Copyright 2024 The HuggingFace Inc. 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.
+from typing import TYPE_CHECKING, Dict, List, Optional, Union
+
+from .base import HfQuantizer
+
+
+if TYPE_CHECKING:
+ from ..modeling_utils import PreTrainedModel
+
+from ..utils import is_accelerate_available, is_torch_available, logging
+
+
+if is_torch_available():
+ import torch
+
+
+logger = logging.get_logger(__name__)
+
+
+class BitNetHfQuantizer(HfQuantizer):
+ """
+ 1.58-bit quantization from BitNet quantization method:
+ Before loading: it converts the linear layers into BitLinear layers during loading.
+
+ Checkout the paper introducing this method : https://arxiv.org/pdf/2402.17764
+ """
+
+ requires_parameters_quantization = False
+ requires_calibration = True
+
+ required_packages = ["accelerate"]
+
+ def __init__(self, quantization_config, **kwargs):
+ super().__init__(quantization_config, **kwargs)
+ self.quantization_config = quantization_config
+
+ def validate_environment(self, *args, **kwargs):
+ if not is_accelerate_available():
+ raise ImportError("Loading a BitNet quantized model requires accelerate (`pip install accelerate`)")
+
+ if kwargs.get("from_tf", False) or kwargs.get("from_flax", False):
+ raise ValueError(
+ "Loading ternary weights from tf/flax is currently not supported, please make"
+ " sure the weights are in PyTorch format."
+ )
+
+ if not torch.cuda.is_available():
+ logger.warning_once(
+ "You don't have a GPU available to load the model, the inference will be slow because of weight unpacking"
+ )
+ return
+
+ device_map = kwargs.get("device_map", None)
+ if device_map is None:
+ logger.warning_once(
+ "You have loaded a BitNet model on CPU and have a CUDA device available, make sure to set "
+ "your model on a GPU device in order to run your model."
+ )
+ elif device_map is not None:
+ if isinstance(device_map, dict) and ("cpu" in device_map.values() or "disk" in device_map.values()):
+ raise ValueError(
+ "You are attempting to load a BitNet model with a device_map that contains a CPU or disk device."
+ "This is not supported. Please remove the CPU or disk device from the device_map."
+ )
+
+ def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs):
+ return model
+
+ def _process_model_before_weight_loading(
+ self,
+ model: "PreTrainedModel",
+ keep_in_fp32_modules: Optional[List[str]] = None,
+ **kwargs,
+ ):
+ from ..integrations import replace_with_bitnet_linear
+
+ self.modules_to_not_convert = self.get_modules_to_not_convert(
+ model, self.quantization_config.modules_to_not_convert, keep_in_fp32_modules
+ )
+
+ model = replace_with_bitnet_linear(
+ model,
+ modules_to_not_convert=self.modules_to_not_convert,
+ quantization_config=self.quantization_config,
+ pre_quantized=self.pre_quantized,
+ )
+
+ def adjust_max_memory(self, max_memory: Dict[str, Union[int, str]]) -> Dict[str, Union[int, str]]:
+ max_memory = {key: val * 0.90 for key, val in max_memory.items()}
+ return max_memory
+
+ def adjust_target_dtype(self, target_dtype: "torch.dtype") -> "torch.dtype":
+ target_dtype = torch.int8
+ return target_dtype
+
+ def is_serializable(self, safe_serialization=None):
+ return True
+
+ @property
+ def is_trainable(self) -> bool:
+ return False
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_bnb_4bit.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_bnb_4bit.py
new file mode 100644
index 0000000000000000000000000000000000000000..7fb9176c4673a009914600ae47068c22826739d2
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_bnb_4bit.py
@@ -0,0 +1,358 @@
+# Copyright 2024 The HuggingFace Inc. 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.
+import importlib
+from functools import cached_property
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+
+from packaging import version
+
+from .base import HfQuantizer
+from .quantizers_utils import get_module_from_name
+
+
+if TYPE_CHECKING:
+ from ..modeling_utils import PreTrainedModel
+
+from ..utils import (
+ ACCELERATE_MIN_VERSION,
+ is_accelerate_available,
+ is_bitsandbytes_available,
+ is_torch_available,
+ is_torch_hpu_available,
+ is_torch_npu_available,
+ is_torch_xpu_available,
+ logging,
+)
+
+
+if is_torch_available():
+ import torch
+
+ from ..pytorch_utils import Conv1D
+
+logger = logging.get_logger(__name__)
+
+
+class Bnb4BitHfQuantizer(HfQuantizer):
+ """
+ 4-bit quantization from bitsandbytes.py quantization method:
+ before loading: converts transformer layers into Linear4bit during loading: load 16bit weight and pass to the
+ layer object after: quantizes individual weights in Linear4bit into 4bit at the first .cuda() call
+ saving:
+ from state dict, as usual; saves weights and `quant_state` components
+ loading:
+ need to locate `quant_state` components and pass to Param4bit constructor
+ """
+
+ use_keep_in_fp32_modules = True
+ requires_parameters_quantization = True
+ requires_calibration = False
+
+ required_packages = ["bitsandbytes", "accelerate"]
+
+ def __init__(self, quantization_config, **kwargs):
+ super().__init__(quantization_config, **kwargs)
+
+ if self.quantization_config.llm_int8_skip_modules is not None:
+ self.modules_to_not_convert = self.quantization_config.llm_int8_skip_modules
+
+ def validate_environment(self, *args, **kwargs):
+ if not is_accelerate_available():
+ raise ImportError(
+ f"Using `bitsandbytes` 4-bit quantization requires Accelerate: `pip install 'accelerate>={ACCELERATE_MIN_VERSION}'`"
+ )
+ if not is_bitsandbytes_available():
+ raise ImportError(
+ "Using `bitsandbytes` 4-bit quantization requires the latest version of bitsandbytes: `pip install -U bitsandbytes`"
+ )
+
+ from ..integrations import validate_bnb_backend_availability
+ from ..utils import is_bitsandbytes_multi_backend_available
+
+ bnb_multibackend_is_enabled = is_bitsandbytes_multi_backend_available()
+ validate_bnb_backend_availability(raise_exception=True)
+
+ if kwargs.get("from_tf", False) or kwargs.get("from_flax", False):
+ raise ValueError(
+ "Converting into 4-bit or 8-bit weights from tf/flax weights is currently not supported, please make"
+ " sure the weights are in PyTorch format."
+ )
+
+ device_map = kwargs.get("device_map", None)
+ if (
+ device_map is not None
+ and isinstance(device_map, dict)
+ and not self.quantization_config.llm_int8_enable_fp32_cpu_offload
+ ):
+ device_map_without_lm_head = {
+ key: device_map[key] for key in device_map.keys() if key not in self.modules_to_not_convert
+ }
+ if set(device_map.values()) == {"cpu"} and bnb_multibackend_is_enabled:
+ pass
+ elif "cpu" in device_map_without_lm_head.values() or "disk" in device_map_without_lm_head.values():
+ raise ValueError(
+ "Some modules are dispatched on the CPU or the disk. Make sure you have enough GPU RAM to fit the "
+ "quantized model. If you want to dispatch the model on the CPU or the disk while keeping these modules "
+ "in 32-bit, you need to set `llm_int8_enable_fp32_cpu_offload=True` and pass a custom `device_map` to "
+ "`from_pretrained`. Check "
+ "https://huggingface.co/docs/transformers/main/en/main_classes/quantization#offload-between-cpu-and-gpu "
+ "for more details. "
+ )
+
+ if version.parse(importlib.metadata.version("bitsandbytes")) < version.parse("0.39.0"):
+ raise ValueError(
+ "You have a version of `bitsandbytes` that is not compatible with 4bit inference and training"
+ " make sure you have the latest version of `bitsandbytes` installed"
+ )
+
+ def adjust_target_dtype(self, target_dtype: "torch.dtype") -> "torch.dtype":
+ if version.parse(importlib.metadata.version("accelerate")) > version.parse("0.19.0"):
+ from accelerate.utils import CustomDtype
+
+ if target_dtype != torch.int8:
+ logger.info("target_dtype {target_dtype} is replaced by `CustomDtype.INT4` for 4-bit BnB quantization")
+ return CustomDtype.INT4
+ else:
+ raise ValueError(
+ "You are using `device_map='auto'` on a 4bit loaded version of the model. To automatically compute"
+ " the appropriate device map, you should upgrade your `accelerate` library,"
+ "`pip install --upgrade accelerate` or install it from source to support fp4 auto device map"
+ "calculation. You may encounter unexpected behavior, or pass your own device map"
+ )
+
+ def check_quantized_param(
+ self,
+ model: "PreTrainedModel",
+ param_value: "torch.Tensor",
+ param_name: str,
+ state_dict: Dict[str, Any],
+ **kwargs,
+ ) -> bool:
+ import bitsandbytes as bnb
+
+ module, tensor_name = get_module_from_name(model, param_name)
+ if isinstance(module._parameters.get(tensor_name, None), bnb.nn.Params4bit):
+ # Add here check for loaded components' dtypes once serialization is implemented
+ return True
+ elif isinstance(module, bnb.nn.Linear4bit) and tensor_name == "bias":
+ # bias could be loaded by regular set_module_tensor_to_device() from accelerate,
+ # but it would wrongly use uninitialized weight there.
+ return True
+ else:
+ return False
+
+ def create_quantized_param(
+ self,
+ model: "PreTrainedModel",
+ param_value: "torch.Tensor",
+ param_name: str,
+ target_device: "torch.device",
+ state_dict: Dict[str, Any],
+ unexpected_keys: Optional[List[str]] = None,
+ ):
+ """
+ combines logic from _load_state_dict_into_meta_model and .integrations.bitsandbytes.py::set_module_quantized_tensor_to_device()
+ """
+ import bitsandbytes as bnb
+
+ module, tensor_name = get_module_from_name(model, param_name)
+
+ if tensor_name not in module._parameters:
+ raise ValueError(f"{module} does not have a parameter or a buffer named {tensor_name}.")
+
+ old_value = getattr(module, tensor_name)
+
+ # `torch.Tensor.to()` is not supported by `torch_npu` (see this [issue](https://github.com/Ascend/pytorch/issues/16)).
+ if isinstance(target_device, int) and is_torch_npu_available():
+ target_device = f"npu:{target_device}"
+ if tensor_name == "bias":
+ if param_value is None:
+ new_value = old_value.to(target_device)
+ else:
+ new_value = param_value.to(target_device)
+
+ new_value = torch.nn.Parameter(new_value, requires_grad=old_value.requires_grad)
+ module._parameters[tensor_name] = new_value
+ return
+
+ if not isinstance(module._parameters[tensor_name], bnb.nn.Params4bit):
+ raise ValueError("this function only loads `Linear4bit components`")
+ if (
+ old_value.device == torch.device("meta")
+ and target_device not in ["meta", torch.device("meta")]
+ and param_value is None
+ ):
+ raise ValueError(f"{tensor_name} is on the meta device, we need a `value` to put in on {target_device}.")
+
+ # construct `new_value` for the module._parameters[tensor_name]:
+ if self.pre_quantized:
+ # 4bit loading. Collecting components for restoring quantized weight
+ # This can be expanded to make a universal call for any quantized weight loading
+
+ if not self.is_serializable:
+ raise ValueError(
+ "Detected int4 weights but the version of bitsandbytes is not compatible with int4 serialization. "
+ "Make sure to download the latest `bitsandbytes` version. `pip install --upgrade bitsandbytes`."
+ )
+
+ if (param_name + ".quant_state.bitsandbytes__fp4" not in state_dict) and (
+ param_name + ".quant_state.bitsandbytes__nf4" not in state_dict
+ ):
+ raise ValueError(
+ f"Supplied state dict for {param_name} does not contain `bitsandbytes__*` and possibly other `quantized_stats` components."
+ )
+
+ quantized_stats = {}
+ for k, v in state_dict.items():
+ if param_name + "." in k:
+ quantized_stats[k] = v
+ if unexpected_keys is not None and k in unexpected_keys:
+ unexpected_keys.remove(k)
+
+ param_kwargs = {}
+ if self.is_bnb_supports_quant_storage_module:
+ param_kwargs["module"] = module
+
+ new_value = bnb.nn.Params4bit.from_prequantized(
+ data=param_value,
+ quantized_stats=quantized_stats,
+ requires_grad=False,
+ device=target_device,
+ **param_kwargs,
+ )
+ else:
+ new_value = param_value.to("cpu")
+
+ # Support models using `Conv1D` in place of `nn.Linear` (e.g. openai-community/gpt2) by transposing the weight matrix prior to quantization.
+ # Since weights are saved in the correct "orientation", we skip transposing when loading.
+ if issubclass(module.source_cls, Conv1D):
+ new_value = new_value.T
+
+ kwargs = old_value.__dict__
+ new_value = bnb.nn.Params4bit(new_value, requires_grad=False, **kwargs).to(target_device)
+
+ module._parameters[tensor_name] = new_value
+
+ # Copied from transformers.quantizers.quantizer_bnb_8bit.Bnb8BitHfQuantizer.adjust_max_memory
+ def adjust_max_memory(self, max_memory: Dict[str, Union[int, str]]) -> Dict[str, Union[int, str]]:
+ # need more space for buffers that are created during quantization
+ max_memory = {key: val * 0.90 for key, val in max_memory.items()}
+ return max_memory
+
+ # Copied from transformers.quantizers.quantizer_bnb_8bit.Bnb8BitHfQuantizer.update_torch_dtype
+ def update_torch_dtype(self, torch_dtype: "torch.dtype") -> "torch.dtype":
+ if torch_dtype is None:
+ # We force the `dtype` to be float16, this is a requirement from `bitsandbytes`
+ logger.info(
+ "Overriding torch_dtype=%s with `torch_dtype=torch.float16` due to "
+ "requirements of `bitsandbytes` to enable model loading in 8-bit or 4-bit. "
+ "Pass your own torch_dtype to specify the dtype of the remaining non-linear layers or pass"
+ " torch_dtype=torch.float16 to remove this warning.",
+ torch_dtype,
+ )
+ torch_dtype = torch.float16
+ return torch_dtype
+
+ def update_device_map(self, device_map):
+ if device_map is None:
+ if torch.cuda.is_available():
+ device_map = {"": torch.cuda.current_device()}
+ elif is_torch_npu_available():
+ device_map = {"": f"npu:{torch.npu.current_device()}"}
+ elif is_torch_hpu_available():
+ device_map = {"": f"hpu:{torch.hpu.current_device()}"}
+ elif is_torch_xpu_available():
+ device_map = {"": f"xpu:{torch.xpu.current_device()}"}
+ else:
+ device_map = {"": "cpu"}
+ logger.info(
+ "The device_map was not initialized. "
+ f"Setting device_map to {device_map}. "
+ "If you want to use the model for inference, please set device_map ='auto' "
+ )
+ return device_map
+
+ # Copied from transformers.quantizers.quantizer_bnb_8bit.Bnb8BitHfQuantizer._process_model_before_weight_loading
+ def _process_model_before_weight_loading(
+ self,
+ model: "PreTrainedModel",
+ device_map,
+ keep_in_fp32_modules: Optional[List[str]] = None,
+ **kwargs,
+ ):
+ from ..integrations import replace_with_bnb_linear
+
+ llm_int8_enable_fp32_cpu_offload = self.quantization_config.llm_int8_enable_fp32_cpu_offload
+
+ self.modules_to_not_convert = self.get_modules_to_not_convert(
+ model, self.quantization_config.llm_int8_skip_modules, keep_in_fp32_modules
+ )
+
+ # Extend `self.modules_to_not_convert` to keys that are supposed to be offloaded to `cpu` or `disk`
+ if isinstance(device_map, dict) and len(device_map.keys()) > 1:
+ keys_on_cpu = [key for key, value in device_map.items() if value in ["disk", "cpu"]]
+
+ if len(keys_on_cpu) > 0 and not llm_int8_enable_fp32_cpu_offload:
+ raise ValueError(
+ "If you want to offload some keys to `cpu` or `disk`, you need to set "
+ "`llm_int8_enable_fp32_cpu_offload=True`. Note that these modules will not be "
+ " converted to 8-bit but kept in 32-bit."
+ )
+ self.modules_to_not_convert.extend(keys_on_cpu)
+
+ model = replace_with_bnb_linear(
+ model, modules_to_not_convert=self.modules_to_not_convert, quantization_config=self.quantization_config
+ )
+ # TODO: consider bringing replace_with_bnb_linear() code from ..integrations/bitsandbyter.py to here
+
+ model.config.quantization_config = self.quantization_config
+
+ # Copied from transformers.quantizers.quantizer_bnb_8bit.Bnb8BitHfQuantizer._process_model_after_weight_loading with 8bit->4bit
+ def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs):
+ model.is_loaded_in_4bit = True
+ model.is_4bit_serializable = self.is_serializable()
+ return model
+
+ def is_serializable(self, safe_serialization=None):
+ _is_4bit_serializable = version.parse(importlib.metadata.version("bitsandbytes")) >= version.parse("0.41.3")
+
+ if not _is_4bit_serializable:
+ logger.warning(
+ "You are calling `save_pretrained` to a 4-bit converted model, but your `bitsandbytes` version doesn't support it. "
+ "If you want to save 4-bit models, make sure to have `bitsandbytes>=0.41.3` installed."
+ )
+ return False
+
+ return True
+
+ @cached_property
+ def is_bnb_supports_quant_storage_module(self) -> bool:
+ """
+ determines if the current version of bitsandbytes supports
+ the `module` parameter in `Params4bit.from_prequantized`
+ :return:
+ """
+ return version.parse(importlib.metadata.version("bitsandbytes")) >= version.parse("0.43.3")
+
+ @property
+ def is_trainable(self) -> bool:
+ return True
+
+ def _dequantize(self, model):
+ from ..integrations import dequantize_and_replace
+
+ model = dequantize_and_replace(
+ model, self.modules_to_not_convert, quantization_config=self.quantization_config
+ )
+ return model
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_bnb_8bit.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_bnb_8bit.py
new file mode 100644
index 0000000000000000000000000000000000000000..cac339b16b9a099f52cc8715eb0b0ecddcff54fd
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_bnb_8bit.py
@@ -0,0 +1,303 @@
+# Copyright 2024 The HuggingFace Inc. 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.
+import importlib
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+
+from packaging import version
+
+from .base import HfQuantizer
+
+
+if TYPE_CHECKING:
+ from ..modeling_utils import PreTrainedModel
+
+from ..utils import (
+ ACCELERATE_MIN_VERSION,
+ is_accelerate_available,
+ is_bitsandbytes_available,
+ is_torch_available,
+ is_torch_xpu_available,
+ logging,
+)
+from .quantizers_utils import get_module_from_name
+
+
+if is_torch_available():
+ import torch
+
+ from ..pytorch_utils import Conv1D
+
+logger = logging.get_logger(__name__)
+
+
+class Bnb8BitHfQuantizer(HfQuantizer):
+ """
+ 8-bit quantization from bitsandbytes quantization method:
+ before loading: converts transformer layers into Linear8bitLt during loading: load 16bit weight and pass to the
+ layer object after: quantizes individual weights in Linear8bitLt into 8bit at fitst .cuda() call
+ saving:
+ from state dict, as usual; saves weights and 'SCB' component
+ loading:
+ need to locate SCB component and pass to the Linear8bitLt object
+ """
+
+ use_keep_in_fp32_modules = True
+ requires_parameters_quantization = True
+ requires_calibration = False
+
+ required_packages = ["bitsandbytes", "accelerate"]
+
+ def __init__(self, quantization_config, **kwargs):
+ super().__init__(quantization_config, **kwargs)
+
+ if self.quantization_config.llm_int8_skip_modules is not None:
+ self.modules_to_not_convert = self.quantization_config.llm_int8_skip_modules
+
+ def validate_environment(self, *args, **kwargs):
+ if not is_accelerate_available():
+ raise ImportError(
+ f"Using `bitsandbytes` 8-bit quantization requires Accelerate: `pip install 'accelerate>={ACCELERATE_MIN_VERSION}'`"
+ )
+ if not is_bitsandbytes_available():
+ raise ImportError(
+ "Using `bitsandbytes` 8-bit quantization requires the latest version of bitsandbytes: `pip install -U bitsandbytes`"
+ )
+
+ from ..integrations import validate_bnb_backend_availability
+ from ..utils import is_bitsandbytes_multi_backend_available
+
+ bnb_multibackend_is_enabled = is_bitsandbytes_multi_backend_available()
+ validate_bnb_backend_availability(raise_exception=True)
+
+ if kwargs.get("from_tf", False) or kwargs.get("from_flax", False):
+ raise ValueError(
+ "Converting into 4-bit or 8-bit weights from tf/flax weights is currently not supported, please make"
+ " sure the weights are in PyTorch format."
+ )
+
+ device_map = kwargs.get("device_map", None)
+ if (
+ device_map is not None
+ and isinstance(device_map, dict)
+ and not self.quantization_config.llm_int8_enable_fp32_cpu_offload
+ ):
+ device_map_without_lm_head = {
+ key: device_map[key] for key in device_map.keys() if key not in self.modules_to_not_convert
+ }
+ if set(device_map.values()) == {"cpu"} and bnb_multibackend_is_enabled:
+ pass
+ elif "cpu" in device_map_without_lm_head.values() or "disk" in device_map_without_lm_head.values():
+ raise ValueError(
+ "Some modules are dispatched on the CPU or the disk. Make sure you have enough GPU RAM to fit the "
+ "quantized model. If you want to dispatch the model on the CPU or the disk while keeping these modules "
+ "in 32-bit, you need to set `llm_int8_enable_fp32_cpu_offload=True` and pass a custom `device_map` to "
+ "`from_pretrained`. Check "
+ "https://huggingface.co/docs/transformers/main/en/main_classes/quantization#offload-between-cpu-and-gpu "
+ "for more details. "
+ )
+
+ if version.parse(importlib.metadata.version("bitsandbytes")) < version.parse("0.37.2"):
+ raise ValueError(
+ "You have a version of `bitsandbytes` that is not compatible with 8bit inference and training"
+ " make sure you have the latest version of `bitsandbytes` installed"
+ )
+
+ def adjust_max_memory(self, max_memory: Dict[str, Union[int, str]]) -> Dict[str, Union[int, str]]:
+ # need more space for buffers that are created during quantization
+ max_memory = {key: val * 0.90 for key, val in max_memory.items()}
+ return max_memory
+
+ def update_torch_dtype(self, torch_dtype: "torch.dtype") -> "torch.dtype":
+ if torch_dtype is None:
+ # We force the `dtype` to be float16, this is a requirement from `bitsandbytes`
+ logger.info(
+ "Overriding torch_dtype=%s with `torch_dtype=torch.float16` due to "
+ "requirements of `bitsandbytes` to enable model loading in 8-bit or 4-bit. "
+ "Pass your own torch_dtype to specify the dtype of the remaining non-linear layers or pass"
+ " torch_dtype=torch.float16 to remove this warning.",
+ torch_dtype,
+ )
+ torch_dtype = torch.float16
+ return torch_dtype
+
+ def update_device_map(self, device_map):
+ if device_map is None:
+ if torch.cuda.is_available():
+ device_map = {"": torch.cuda.current_device()}
+ elif is_torch_xpu_available():
+ device_map = {"": f"xpu:{torch.xpu.current_device()}"}
+ else:
+ device_map = {"": "cpu"}
+ logger.info(
+ "The device_map was not initialized. "
+ f"Setting device_map to {device_map}. "
+ "If you want to use the model for inference, please set device_map ='auto' "
+ )
+ return device_map
+
+ def adjust_target_dtype(self, target_dtype: "torch.dtype") -> "torch.dtype":
+ if target_dtype != torch.int8:
+ logger.info("target_dtype {target_dtype} is replaced by `torch.int8` for 8-bit BnB quantization")
+ return torch.int8
+
+ def check_quantized_param(
+ self,
+ model: "PreTrainedModel",
+ param_value: "torch.Tensor",
+ param_name: str,
+ state_dict: Dict[str, Any],
+ **kwargs,
+ ):
+ import bitsandbytes as bnb
+
+ module, tensor_name = get_module_from_name(model, param_name)
+ if isinstance(module._parameters.get(tensor_name, None), bnb.nn.Int8Params):
+ if self.pre_quantized:
+ if param_name.replace("weight", "SCB") not in state_dict.keys():
+ raise ValueError("Missing quantization component `SCB`")
+ if param_value.dtype != torch.int8:
+ raise ValueError(
+ f"Incompatible dtype `{param_value.dtype}` when loading 8-bit prequantized weight. Expected `torch.int8`."
+ )
+ return True
+ return False
+
+ def create_quantized_param(
+ self,
+ model: "PreTrainedModel",
+ param_value: "torch.Tensor",
+ param_name: str,
+ target_device: "torch.device",
+ state_dict: Dict[str, Any],
+ unexpected_keys: Optional[List[str]] = None,
+ ):
+ """
+ combines logic from _load_state_dict_into_meta_model and .integrations.bitsandbytes.py::set_module_quantized_tensor_to_device()
+ needs aux items from state dicts, if found - removes them from unexpected_keys
+ """
+ import bitsandbytes as bnb
+
+ fp16_statistics_key = param_name.replace("weight", "SCB")
+ fp16_weights_format_key = param_name.replace("weight", "weight_format")
+
+ fp16_statistics = state_dict.get(fp16_statistics_key, None)
+ fp16_weights_format = state_dict.get(fp16_weights_format_key, None)
+
+ module, tensor_name = get_module_from_name(model, param_name)
+ if tensor_name not in module._parameters:
+ raise ValueError(f"{module} does not have a parameter or a buffer named {tensor_name}.")
+
+ old_value = getattr(module, tensor_name)
+
+ if not isinstance(module._parameters[tensor_name], bnb.nn.Int8Params):
+ raise ValueError(f"Parameter `{tensor_name}` should only be a `bnb.nn.Int8Params` instance.")
+ if (
+ old_value.device == torch.device("meta")
+ and target_device not in ["meta", torch.device("meta")]
+ and param_value is None
+ ):
+ raise ValueError(f"{tensor_name} is on the meta device, we need a `value` to put in on {target_device}.")
+
+ new_value = param_value.to("cpu")
+ if self.pre_quantized and not self.is_serializable():
+ raise ValueError(
+ "Detected int8 weights but the version of bitsandbytes is not compatible with int8 serialization. "
+ "Make sure to download the latest `bitsandbytes` version. `pip install --upgrade bitsandbytes`."
+ )
+
+ # Support models using `Conv1D` in place of `nn.Linear` (e.g. openai-community/gpt2) by transposing the weight matrix prior to quantization.
+ # Since weights are saved in the correct "orientation", we skip transposing when loading.
+ if issubclass(module.source_cls, Conv1D):
+ if fp16_statistics is None:
+ new_value = new_value.T
+
+ kwargs = old_value.__dict__
+ new_value = bnb.nn.Int8Params(new_value, requires_grad=False, **kwargs).to(target_device)
+
+ module._parameters[tensor_name] = new_value
+ if fp16_statistics is not None:
+ setattr(module.weight, "SCB", fp16_statistics.to(target_device))
+ if unexpected_keys is not None:
+ unexpected_keys.remove(fp16_statistics_key)
+
+ # We just need to pop the `weight_format` keys from the state dict to remove unneeded
+ # messages. The correct format is correctly retrieved during the first forward pass.
+ if fp16_weights_format is not None and unexpected_keys is not None:
+ unexpected_keys.remove(fp16_weights_format_key)
+
+ def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs):
+ model.is_loaded_in_8bit = True
+ model.is_8bit_serializable = self.is_serializable()
+ return model
+
+ def _process_model_before_weight_loading(
+ self,
+ model: "PreTrainedModel",
+ device_map,
+ keep_in_fp32_modules: Optional[List[str]] = None,
+ **kwargs,
+ ):
+ from ..integrations import replace_with_bnb_linear
+
+ llm_int8_enable_fp32_cpu_offload = self.quantization_config.llm_int8_enable_fp32_cpu_offload
+
+ self.modules_to_not_convert = self.get_modules_to_not_convert(
+ model, self.quantization_config.llm_int8_skip_modules, keep_in_fp32_modules
+ )
+
+ # Extend `self.modules_to_not_convert` to keys that are supposed to be offloaded to `cpu` or `disk`
+ if isinstance(device_map, dict) and len(device_map.keys()) > 1:
+ keys_on_cpu = [key for key, value in device_map.items() if value in ["disk", "cpu"]]
+
+ if len(keys_on_cpu) > 0 and not llm_int8_enable_fp32_cpu_offload:
+ raise ValueError(
+ "If you want to offload some keys to `cpu` or `disk`, you need to set "
+ "`llm_int8_enable_fp32_cpu_offload=True`. Note that these modules will not be "
+ " converted to 8-bit but kept in 32-bit."
+ )
+ self.modules_to_not_convert.extend(keys_on_cpu)
+
+ model = replace_with_bnb_linear(
+ model, modules_to_not_convert=self.modules_to_not_convert, quantization_config=self.quantization_config
+ )
+ # TODO: consider bringing replace_with_bnb_linear() code from ..integrations/bitsandbyter.py to here
+
+ model.config.quantization_config = self.quantization_config
+
+ def is_serializable(self, safe_serialization=None):
+ _bnb_supports_8bit_serialization = version.parse(importlib.metadata.version("bitsandbytes")) > version.parse(
+ "0.37.2"
+ )
+
+ if not _bnb_supports_8bit_serialization:
+ logger.warning(
+ "You are calling `save_pretrained` to a 8-bit converted model, but your `bitsandbytes` version doesn't support it. "
+ "If you want to save 8-bit models, make sure to have `bitsandbytes>0.37.2` installed. You will most likely face errors or"
+ " unexpected behaviours."
+ )
+ return False
+
+ return True
+
+ @property
+ def is_trainable(self) -> bool:
+ return version.parse(importlib.metadata.version("bitsandbytes")) >= version.parse("0.37.0")
+
+ def _dequantize(self, model):
+ from ..integrations import dequantize_and_replace
+
+ model = dequantize_and_replace(
+ model, self.modules_to_not_convert, quantization_config=self.quantization_config
+ )
+ return model
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_compressed_tensors.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_compressed_tensors.py
new file mode 100644
index 0000000000000000000000000000000000000000..4e45abf95340b51138983a3a7132ab2f98b295de
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_compressed_tensors.py
@@ -0,0 +1,173 @@
+# Copyright 2024 The HuggingFace Inc. 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.
+
+
+import os
+import re
+from typing import List
+
+from ..utils import is_compressed_tensors_available, is_torch_available, logging
+from ..utils.quantization_config import CompressedTensorsConfig
+from .base import HfQuantizer
+
+
+if is_torch_available():
+ import torch
+
+logger = logging.get_logger(__name__)
+
+
+class CompressedTensorsHfQuantizer(HfQuantizer):
+ """
+ Quantizer for the compressed_tensors package. Loads and restores models to
+ quantized state with compressed_tensors
+ """
+
+ requires_calibration = True
+ required_packages = ["compressed_tensors"]
+
+ def __init__(self, quantization_config: CompressedTensorsConfig, **kwargs):
+ super().__init__(quantization_config, **kwargs)
+
+ if not is_compressed_tensors_available():
+ raise ImportError(
+ "Using `compressed_tensors` quantized models requires the compressed-tensors library: "
+ "`pip install compressed-tensors`"
+ )
+
+ # Call post_init here to ensure proper config setup when `run_compressed`
+ # is provided directly via CompressedTensorsConfig, and to avoid duplicate logging.
+
+ quantization_config.post_init()
+ from compressed_tensors.compressors import ModelCompressor
+
+ self.compressor = ModelCompressor.from_compression_config(quantization_config)
+ self.run_compressed = quantization_config.run_compressed
+ self.quantization_config = quantization_config
+
+ def update_missing_keys_after_loading(self, model, missing_keys: List[str], prefix: str) -> List[str]:
+ """
+ Update missing keys after loading the model. This is necessary for compressed tensors
+ to load the model correctly. We expect weights to be present in missing keys.
+ The weight's are re-constructed by ModelCompressor in _process_model_after_weight_loading
+
+ This function cleans up expected missing keys and returns the remaining missing keys
+ """
+
+ if self.run_compressed:
+ return missing_keys
+
+ # We expect some keys to be missing for
+ # compresed models
+ # This is fine as the weights are reconstructed by ModelCompressor
+ # in _process_model_after_weight_loading
+
+ expected_missing_keys = self.compressor.get_missing_module_keys(model)
+ return [
+ key for key in missing_keys if not any(re.match(f".*{pattern}", key) for pattern in expected_missing_keys)
+ ]
+
+ def update_unexpected_keys(self, model, unexpected_keys: List[str], prefix: str) -> List[str]:
+ """
+ Override this method if you want to adjust the `unexpected_keys`.
+
+ Args:
+ unexpected_keys (`List[str]`, *optional*):
+ The list of unexpected keys in the checkpoint compared to the state dict of the model
+ """
+
+ if self.run_compressed:
+ return unexpected_keys
+
+ # We expect some unexpected keys in model
+ # safetensors file for compressed models
+ keys_to_ignore = self.compressor.get_unexpected_file_keys(model)
+ return [key for key in unexpected_keys if not any(re.match(f".*{pattern}", key) for pattern in keys_to_ignore)]
+
+ def validate_environment(self, *args, **kwargs):
+ if not is_compressed_tensors_available():
+ raise ImportError(
+ "Using `compressed_tensors` quantized models requires the compressed-tensors library: "
+ "`pip install compressed-tensors`"
+ )
+ if not is_torch_available():
+ # torch already should be installed as part of compressed tensors
+ raise ImportError("torch is required for using compressed-tensors quantization")
+
+ def update_torch_dtype(self, torch_dtype: "torch.dtype") -> "torch.dtype":
+ if torch_dtype is None:
+ logger.info("Loading model using torch.float16 for compressed-tensors quantization")
+ torch_dtype = torch.float16
+ elif torch_dtype != torch.float16:
+ logger.info(
+ "We suggest you to set `torch_dtype=torch.float16` for better efficiency with compressed_tensors."
+ )
+ return torch_dtype
+
+ def _process_model_before_weight_loading(self, model, **kwargs):
+ from compressed_tensors.quantization import apply_quantization_config
+
+ ct_quantization_config = self.compressor.quantization_config
+
+ if self.run_compressed:
+ apply_quantization_config(model, ct_quantization_config, run_compressed=True)
+ elif not self.quantization_config.is_quantization_compressed:
+ apply_quantization_config(model, ct_quantization_config)
+
+ def _process_model_after_weight_loading(self, model, **kwargs):
+ """Decompress loaded model if necessary - need for qat"""
+
+ if (
+ self.quantization_config.is_quantization_compressed and not self.run_compressed
+ ) or self.quantization_config.is_sparsification_compressed:
+ config = kwargs.get("config", None)
+ cache_path = config._name_or_path
+
+ if not os.path.exists(cache_path):
+ from transformers.utils import cached_file
+
+ config_file_path = cached_file(cache_path, "config.json")
+ cache_path = os.path.sep.join(config_file_path.split(os.path.sep)[:-1])
+
+ if self.quantization_config.is_quantization_compressed and not self.run_compressed:
+ from compressed_tensors.quantization import QuantizationStatus
+
+ self.compressor.quantization_config.quantization_status = QuantizationStatus.FROZEN
+ self.compressor.decompress(model_path=cache_path, model=model)
+
+ def update_tp_plan(self, config):
+ additional_plan = {
+ "layers.*.feed_forward.experts.*.gate_proj.weight": "local_colwise",
+ "layers.*.feed_forward.experts.*.gate_proj.weight_scale": "local_colwise",
+ "layers.*.feed_forward.experts.*.up_proj.weight": "local_colwise",
+ "layers.*.feed_forward.experts.*.up_proj.weight_scale": "local_colwise",
+ "layers.*.feed_forward.experts.*.down_proj.weight": "local_rowwise",
+ }
+ if config.get_text_config() is not None and config.get_text_config().base_model_tp_plan is not None:
+ config.get_text_config().base_model_tp_plan.update(additional_plan)
+
+ return config
+
+ @property
+ def is_trainable(self):
+ return True
+
+ def is_qat_trainable(self) -> bool:
+ """Loaded Models can carry out quantization aware training"""
+ # models need to be decompressed carry out qat
+ return not self.run_compressed or not self.quantization_config.is_quantization_compressed
+
+ def is_serializable(self, safe_serialization=None) -> bool:
+ """Models quantized using compressed tensors can be saved to disk"""
+ return True
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_eetq.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_eetq.py
new file mode 100644
index 0000000000000000000000000000000000000000..988f90789ac3751ca27a9ce4d1ed71762fd7c434
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_eetq.py
@@ -0,0 +1,181 @@
+# Copyright 2024 The HuggingFace Inc. 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.
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
+
+from .base import HfQuantizer
+
+
+if TYPE_CHECKING:
+ from ..modeling_utils import PreTrainedModel
+
+from ..utils import is_accelerate_available, is_eetq_available, is_torch_available, logging
+from .quantizers_utils import get_module_from_name
+
+
+if is_torch_available():
+ import torch
+
+
+logger = logging.get_logger(__name__)
+
+
+class EetqHfQuantizer(HfQuantizer):
+ """
+ 8-bit quantization from EETQ quantization method:
+ before loading: converts transformer layers into W8A16Linear during loading: load 16bit weight and pass to the
+ layer object after: quantizes individual weights in Linear8bitLt into 8bit at first .cuda() call
+ """
+
+ requires_parameters_quantization = True
+ requires_calibration = False
+
+ required_packages = ["eetq", "accelerate"]
+
+ def __init__(self, quantization_config, **kwargs):
+ super().__init__(quantization_config, **kwargs)
+ self.quantization_config = quantization_config
+
+ def validate_environment(self, *args, **kwargs):
+ if not is_eetq_available():
+ raise ImportError(
+ "Using `eetq` 8-bit quantization requires eetq."
+ "Please install the latest version of eetq from : https://github.com/NetEase-FuXi/EETQ"
+ )
+
+ try:
+ import eetq # noqa: F401
+ except ImportError as exc:
+ if "shard_checkpoint" in str(exc):
+ # EETQ 1.0.0 is currently broken with the latest transformers because it tries to import the removed
+ # shard_checkpoint function, see https://github.com/NetEase-FuXi/EETQ/issues/34.
+ # TODO: Update message once eetq releases a fix
+ raise ImportError(
+ "You are using a version of EETQ that is incompatible with the current transformers version. "
+ "Either downgrade transformers to <= v4.46.3 or, if available, upgrade EETQ to > v1.0.0."
+ ) from exc
+ else:
+ raise
+
+ if not is_accelerate_available():
+ raise ImportError("Loading an EETQ quantized model requires accelerate (`pip install accelerate`)")
+
+ if kwargs.get("from_tf", False) or kwargs.get("from_flax", False):
+ raise ValueError(
+ "Converting into 8-bit weights from tf/flax weights is currently not supported, please make"
+ " sure the weights are in PyTorch format."
+ )
+
+ if not torch.cuda.is_available():
+ raise RuntimeError("No GPU found. A GPU is needed for quantization.")
+
+ device_map = kwargs.get("device_map", None)
+ if device_map is None:
+ logger.warning_once(
+ "You have loaded an EETQ model on CPU and have a CUDA device available, make sure to set "
+ "your model on a GPU device in order to run your model."
+ )
+ elif device_map is not None:
+ if isinstance(device_map, dict) and ("cpu" in device_map.values() or "disk" in device_map.values()):
+ raise ValueError(
+ "You are attempting to load an EETQ model with a device_map that contains a CPU or disk device."
+ " This is not supported. Please remove the CPU or disk device from the device_map."
+ )
+
+ def update_torch_dtype(self, torch_dtype: "torch.dtype") -> "torch.dtype":
+ if torch_dtype is None:
+ torch_dtype = torch.float16
+ logger.info(
+ "Overriding torch_dtype=%s with `torch_dtype=torch.float16` due to "
+ "requirements of `eetq` to enable model loading in 8-bit. "
+ "Pass your own torch_dtype to specify the dtype of the remaining non-linear layers or pass"
+ " torch_dtype=torch.float16 to remove this warning.",
+ torch_dtype,
+ )
+ elif torch_dtype != torch.float16:
+ logger.info("We suggest you to set `torch_dtype=torch.float16` for better efficiency with EETQ.")
+ return torch_dtype
+
+ def check_quantized_param(
+ self,
+ model: "PreTrainedModel",
+ param_value: "torch.Tensor",
+ param_name: str,
+ state_dict: Dict[str, Any],
+ **kwargs,
+ ):
+ from eetq import EetqLinear
+
+ module, tensor_name = get_module_from_name(model, param_name)
+
+ if isinstance(module, EetqLinear):
+ if self.pre_quantized or tensor_name == "bias":
+ if tensor_name == "weight" and param_value.dtype != torch.int8:
+ raise ValueError("Expect quantized weights but got an unquantized weight")
+ return False
+ else:
+ if tensor_name == "weight_scale":
+ raise ValueError("Expect unquantized weights but got a quantized weight_scale")
+ return True
+ return False
+
+ def create_quantized_param(
+ self,
+ model: "PreTrainedModel",
+ param_value: "torch.Tensor",
+ param_name: str,
+ target_device: "torch.device",
+ state_dict: Dict[str, Any],
+ unexpected_keys: Optional[List[str]] = None,
+ ):
+ """
+ quantizes weights into qweight and weight_scales
+ """
+ from eetq import quantize_and_preprocess_weights
+
+ module, tensor_name = get_module_from_name(model, param_name)
+ new_value, weight_scale = quantize_and_preprocess_weights(param_value)
+
+ module._buffers[tensor_name] = new_value.to(target_device)
+ module.register("weight_scales", weight_scale.to(target_device))
+
+ def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs):
+ return model
+
+ def _process_model_before_weight_loading(
+ self,
+ model: "PreTrainedModel",
+ keep_in_fp32_modules: Optional[List[str]] = None,
+ **kwargs,
+ ):
+ from ..integrations import replace_with_eetq_linear
+
+ self.modules_to_not_convert = self.get_modules_to_not_convert(
+ model, self.quantization_config.modules_to_not_convert, keep_in_fp32_modules
+ )
+
+ model = replace_with_eetq_linear(
+ model,
+ modules_to_not_convert=self.modules_to_not_convert,
+ quantization_config=self.quantization_config,
+ pre_quantized=self.pre_quantized,
+ )
+
+ model.config.quantization_config = self.quantization_config
+
+ def is_serializable(self, safe_serialization=None):
+ return True
+
+ @property
+ def is_trainable(self) -> bool:
+ return True
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_fbgemm_fp8.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_fbgemm_fp8.py
new file mode 100644
index 0000000000000000000000000000000000000000..074b9fbb8424601350e332c76e9c49f34a737c11
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_fbgemm_fp8.py
@@ -0,0 +1,288 @@
+# Copyright 2024 The HuggingFace Inc. 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.
+import importlib
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
+
+from packaging import version
+
+from .base import HfQuantizer
+
+
+if TYPE_CHECKING:
+ from ..modeling_utils import PreTrainedModel
+
+from ..utils import is_accelerate_available, is_fbgemm_gpu_available, is_torch_available, logging
+from .quantizers_utils import get_module_from_name
+
+
+if is_torch_available():
+ import torch
+
+
+logger = logging.get_logger(__name__)
+
+
+class FbgemmFp8HfQuantizer(HfQuantizer):
+ """
+ FP8 quantization using fbgemm kernels
+ """
+
+ requires_parameters_quantization = True
+ requires_calibration = False
+
+ required_packages = ["fbgemm-gpu", "accelerate"]
+
+ def __init__(self, quantization_config, **kwargs):
+ super().__init__(quantization_config, **kwargs)
+ self.quantization_config = quantization_config
+
+ def validate_environment(self, *args, **kwargs):
+ if not is_torch_available() or version.parse(importlib.metadata.version("torch")) < version.parse("2.1.0"):
+ raise ImportError(
+ "Using fbgemm fp8 quantization requires torch > 2.1.0"
+ "Please install the latest version of torch ( pip install --upgrade torch )"
+ )
+ if not is_fbgemm_gpu_available():
+ raise ImportError(
+ "Using fbgemm fp8 quantization requires fbgemm-gpu library"
+ "Please install the latest version of fbgemm-gpu library by following : https://pytorch.org/FBGEMM/fbgemm_gpu-development/InstallationInstructions.html#fbgemm-gpu-install-libraries"
+ )
+
+ if not is_accelerate_available("0.32.2"):
+ raise ImportError(
+ "Loading an FP8 quantized model requires accelerate > 0.32.1 (`pip install --upgrade accelerate`)"
+ )
+
+ if not torch.cuda.is_available():
+ raise RuntimeError("Using FP8 quantized models with fbgemm kernels requires a GPU")
+
+ compute_capability = torch.cuda.get_device_capability()
+ major, minor = compute_capability
+ if major < 9:
+ raise ValueError(
+ "FP8 quantized models is only supported on GPUs with compute capability >= 9.0 (e.g H100)"
+ )
+
+ device_map = kwargs.get("device_map", None)
+ if device_map is None:
+ logger.warning_once(
+ "You have loaded an FP8 model on CPU and have a CUDA device available, make sure to set "
+ "your model on a GPU device in order to run your model. To remove this warning, pass device_map = 'cuda'. "
+ )
+ elif device_map is not None:
+ if (
+ not self.pre_quantized
+ and isinstance(device_map, dict)
+ and ("cpu" in device_map.values() or "disk" in device_map.values())
+ ):
+ raise ValueError(
+ "You are attempting to load an FP8 model with a device_map that contains a CPU or disk device."
+ "This is not supported when the model is quantized on the fly. "
+ "Please use a quantized checkpoint or remove the CPU or disk device from the device_map."
+ )
+
+ def update_torch_dtype(self, torch_dtype: "torch.dtype") -> "torch.dtype":
+ if torch_dtype is None:
+ torch_dtype = torch.bfloat16
+ logger.info(
+ "Overriding torch_dtype=%s with `torch_dtype=torch.bloat16` due to "
+ "requirements of `fbgemm-gpu` to enable model loading in fp8. "
+ "Pass your own torch_dtype to specify the dtype of the remaining non-linear layers or pass"
+ " torch_dtype=torch.bfloat16 to remove this warning.",
+ torch_dtype,
+ )
+ elif torch_dtype == torch.float16:
+ raise ValueError(
+ "You cannot use FP8 with torch_dtype=torch.float16.We recommend you passing torch_dtype=torch.bfloat16"
+ )
+ return torch_dtype
+
+ def check_quantized_param(
+ self,
+ model: "PreTrainedModel",
+ param_value: "torch.Tensor",
+ param_name: str,
+ state_dict: Dict[str, Any],
+ **kwargs,
+ ):
+ from ..integrations import FbgemmFp8Linear, FbgemmFp8Llama4TextExperts
+
+ module, tensor_name = get_module_from_name(model, param_name)
+
+ if isinstance(module, FbgemmFp8Linear):
+ if self.pre_quantized or tensor_name == "bias":
+ if tensor_name == "weight" and param_value.dtype != torch.float8_e4m3fn:
+ raise ValueError("Expect quantized weights but got an unquantized weight")
+ return False
+ else:
+ if tensor_name == "weight_scale":
+ raise ValueError("Expect unquantized weights but got a quantized weight_scale")
+ return True
+ if isinstance(module, FbgemmFp8Llama4TextExperts):
+ if self.pre_quantized or tensor_name == "bias":
+ return False
+ else:
+ if tensor_name == "gate_up_proj_scale" or tensor_name == "down_proj_scale":
+ raise ValueError("Expect unquantized weights but got a quantized weight_scale")
+ return True
+ return False
+
+ def create_quantized_param(
+ self,
+ model: "PreTrainedModel",
+ param_value: "torch.Tensor",
+ param_name: str,
+ target_device: "torch.device",
+ state_dict: Dict[str, Any],
+ unexpected_keys: Optional[List[str]] = None,
+ ):
+ """
+ Quantizes weights into weight and weight_scale
+ """
+
+ from ..integrations import FbgemmFp8Llama4TextExperts
+
+ module, tensor_name = get_module_from_name(model, param_name)
+ if isinstance(module, FbgemmFp8Llama4TextExperts):
+ if tensor_name == "gate_up_proj":
+ # Process each expert separately
+ # Transpose the second and third dimension
+ transposed_param = param_value.transpose(1, 2)
+
+ # Reshape to 2D for quantization
+ original_shape = transposed_param.shape
+ flattened_param = transposed_param.reshape(-1, original_shape[-1])
+
+ # Quantize using per row instead of per column
+ new_value_flat, weight_scale_flat = torch.ops.fbgemm.quantize_fp8_per_row(flattened_param)
+
+ # Reshape back to original dimensions
+ new_value = new_value_flat.reshape(original_shape)
+ new_value = new_value.transpose(1, 2)
+ weight_scale = weight_scale_flat.reshape(original_shape[0], 1, original_shape[1])
+ elif tensor_name == "down_proj":
+ # Process each expert separately
+ # Transpose the weights for proper quantization
+ transposed_param = param_value.transpose(1, 2)
+
+ # Reshape to 2D for quantization
+ original_shape = transposed_param.shape
+ flattened_param = transposed_param.reshape(-1, original_shape[-1])
+
+ # Quantize using per column
+ new_value_flat, weight_scale_flat = torch.ops.fbgemm.quantize_fp8_per_row(flattened_param)
+
+ # Reshape back to original dimensions
+ new_value = new_value_flat.reshape(original_shape)
+ new_value = new_value.transpose(1, 2)
+ weight_scale = weight_scale_flat.reshape(original_shape[0], original_shape[1], 1)
+
+ module._parameters[f"{tensor_name}_scale"] = torch.nn.Parameter(weight_scale.to(target_device))
+ else:
+ new_value, weight_scale = torch.ops.fbgemm.quantize_fp8_per_row(param_value)
+ module._parameters[f"{tensor_name}_scale"] = torch.nn.Parameter(
+ weight_scale.view(weight_scale.shape[0], 1).to(target_device)
+ )
+
+ module._parameters[tensor_name] = torch.nn.Parameter(new_value.to(target_device))
+
+ if unexpected_keys is not None and param_name in unexpected_keys:
+ unexpected_keys.remove(param_name)
+ del param_name
+
+ def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs):
+ return model
+
+ def _process_model_before_weight_loading(
+ self,
+ model: "PreTrainedModel",
+ keep_in_fp32_modules: Optional[List[str]] = None,
+ **kwargs,
+ ):
+ from ..integrations import replace_with_fbgemm_fp8_linear
+
+ tp_plan = model._tp_plan
+ self.modules_to_not_convert = self.get_modules_to_not_convert(
+ model, self.quantization_config.modules_to_not_convert, keep_in_fp32_modules
+ )
+
+ config = model.config
+ model = replace_with_fbgemm_fp8_linear(
+ model,
+ modules_to_not_convert=self.modules_to_not_convert,
+ quantization_config=self.quantization_config,
+ pre_quantized=self.pre_quantized,
+ config=config,
+ tp_plan=tp_plan,
+ )
+
+ model.config.quantization_config = self.quantization_config
+
+ def update_missing_keys(self, model, missing_keys: List[str], prefix: str) -> List[str]:
+ from ..integrations import FbgemmFp8Linear, FbgemmFp8Llama4TextExperts
+
+ not_missing_keys = []
+ for name, module in model.named_modules():
+ if isinstance(module, FbgemmFp8Linear) or isinstance(module, FbgemmFp8Llama4TextExperts):
+ for missing in missing_keys:
+ if (
+ (name in missing or name in f"{prefix}.{missing}")
+ and not missing.endswith(".weight")
+ and not missing.endswith(".bias")
+ ):
+ not_missing_keys.append(missing)
+ return [k for k in missing_keys if k not in not_missing_keys]
+
+ def update_tp_plan(self, config):
+ text_plan = {
+ "layers.*.self_attn.q_proj.weight": "local_colwise",
+ "layers.*.self_attn.q_proj.weight_scale": "local_colwise",
+ "layers.*.self_attn.k_proj.weight": "local_colwise",
+ "layers.*.self_attn.k_proj.weight_scale": "local_colwise",
+ "layers.*.self_attn.v_proj.weight": "local_colwise",
+ "layers.*.self_attn.v_proj.weight_scale": "local_colwise",
+ "layers.*.self_attn.o_proj.weight": "local_rowwise",
+ "layers.*.self_attn": "gather",
+ "layers.*.input_layernorm.weight": "sequence_parallel",
+ "layers.*.post_attention_layernorm.weight": "sequence_parallel",
+ "norm.weight": "sequence_parallel",
+ "layers.*.feed_forward.shared_expert.gate_proj.weight": "local_colwise",
+ "layers.*.feed_forward.shared_expert.gate_proj.weight_scale": "local_colwise",
+ "layers.*.feed_forward.shared_expert.up_proj.weight": "local_colwise",
+ "layers.*.feed_forward.shared_expert.up_proj.weight_scale": "local_colwise",
+ "layers.*.feed_forward.shared_expert.down_proj.weight": "local_rowwise",
+ "layers.*.feed_forward.experts": "local",
+ "layers.*.feed_forward": "gather",
+ "layers.*.feed_forward.experts.*.gate_proj.weight": "local_colwise",
+ "layers.*.feed_forward.experts.*.gate_proj.weight_scale": "local_colwise",
+ "layers.*.feed_forward.experts.*.up_proj.weight": "local_colwise",
+ "layers.*.feed_forward.experts.*.up_proj.weight_scale": "local_colwise",
+ "layers.*.feed_forward.experts.*.down_proj.weight": "local_rowwise",
+ # For Fused implementation
+ "layers.*.feed_forward.experts.gate_up_proj": "local_packed_rowwise",
+ "layers.*.feed_forward.experts.gate_up_proj_scale": "local_packed_rowwise",
+ "layers.*.feed_forward.experts.down_proj": "local_colwise",
+ }
+ if config.get_text_config() is not None:
+ config.get_text_config().base_model_tp_plan = text_plan
+ else:
+ config.base_model_tp_plan = text_plan
+ return config
+
+ def is_serializable(self, safe_serialization=None):
+ return True
+
+ @property
+ def is_trainable(self) -> bool:
+ return False
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_finegrained_fp8.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_finegrained_fp8.py
new file mode 100644
index 0000000000000000000000000000000000000000..16ce7f6a9eb633959afbf2bcac0a4c18378304f4
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_finegrained_fp8.py
@@ -0,0 +1,206 @@
+import importlib
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
+
+from packaging import version
+
+from ..utils import is_accelerate_available, is_torch_available, logging
+from .base import HfQuantizer
+from .quantizers_utils import get_module_from_name
+
+
+if is_torch_available():
+ import torch
+
+if TYPE_CHECKING:
+ from ..modeling_utils import PreTrainedModel
+
+logger = logging.get_logger(__name__)
+
+
+class FineGrainedFP8HfQuantizer(HfQuantizer):
+ """
+ FP8 quantization implementation supporting both standard and MoE models.
+ Supports both e4m3fn formats based on platform.
+ """
+
+ requires_parameters_quantization = True
+ requires_calibration = False
+ required_packages = ["accelerate"]
+
+ def __init__(self, quantization_config, **kwargs):
+ super().__init__(quantization_config, **kwargs)
+ self.quantization_config = quantization_config
+
+ def validate_environment(self, *args, **kwargs):
+ if not is_torch_available() or version.parse(importlib.metadata.version("torch")) < version.parse("2.1.0"):
+ raise ImportError(
+ "Using fp8 quantization requires torch >= 2.1.0"
+ "Please install the latest version of torch ( pip install --upgrade torch )"
+ )
+
+ if not is_accelerate_available():
+ raise ImportError("Loading an FP8 quantized model requires accelerate (`pip install accelerate`)")
+
+ if kwargs.get("from_tf", False) or kwargs.get("from_flax", False):
+ raise ValueError(
+ "Converting into FP8 weights from tf/flax weights is currently not supported, "
+ "please make sure the weights are in PyTorch format."
+ )
+
+ if not torch.cuda.is_available():
+ raise RuntimeError("No GPU found. A GPU is needed for FP8 quantization.")
+
+ compute_capability = torch.cuda.get_device_capability()
+ major, minor = compute_capability
+ if (major < 8) or (major == 8 and minor < 9):
+ raise ValueError(
+ "FP8 quantized models is only supported on GPUs with compute capability >= 8.9 (e.g 4090/H100)"
+ f", actual = `{major}.{minor}`"
+ )
+
+ device_map = kwargs.get("device_map", None)
+ if device_map is None:
+ logger.warning_once(
+ "You have loaded an FP8 model on CPU and have a CUDA device available, make sure to set "
+ "your model on a GPU device in order to run your model. To remove this warning, pass device_map = 'cuda'. "
+ )
+ elif device_map is not None:
+ if (
+ not self.pre_quantized
+ and isinstance(device_map, dict)
+ and ("cpu" in device_map.values() or "disk" in device_map.values())
+ ):
+ raise ValueError(
+ "You are attempting to load an FP8 model with a device_map that contains a cpu/disk device."
+ "This is not supported when the model is quantized on the fly. "
+ "Please use a quantized checkpoint or remove the cpu/disk device from the device_map."
+ )
+
+ def update_torch_dtype(self, torch_dtype: "torch.dtype") -> "torch.dtype":
+ if torch_dtype is None:
+ logger.info("Setting torch_dtype to torch.float32 as no torch_dtype was specified in from_pretrained")
+ torch_dtype = torch.float32
+ return torch_dtype
+
+ def create_quantized_param(
+ self,
+ model: "PreTrainedModel",
+ param_value: "torch.Tensor",
+ param_name: str,
+ target_device: "torch.device",
+ state_dict: Dict[str, Any],
+ unexpected_keys: Optional[List[str]] = None,
+ ):
+ """
+ Quantizes weights to FP8 format using Block-wise quantization
+ """
+ from accelerate.utils import set_module_tensor_to_device
+
+ set_module_tensor_to_device(model, param_name, target_device, param_value)
+
+ module, tensor_name = get_module_from_name(model, param_name)
+
+ # Get FP8 min/max values
+ fp8_min = torch.finfo(torch.float8_e4m3fn).min
+ fp8_max = torch.finfo(torch.float8_e4m3fn).max
+
+ block_size_m, block_size_n = self.quantization_config.weight_block_size
+
+ rows, cols = param_value.shape[-2:]
+
+ if rows % block_size_m != 0 or cols % block_size_n != 0:
+ raise ValueError(
+ f"Matrix dimensions ({rows}, {cols}) must be divisible by block sizes ({block_size_m}, {block_size_n})"
+ )
+ param_value_orig_shape = param_value.shape
+
+ param_value = param_value.reshape(
+ -1, rows // block_size_m, block_size_m, cols // block_size_n, block_size_n
+ ).permute(0, 1, 3, 2, 4)
+
+ # Calculate scaling factor for each block
+ max_abs = torch.amax(torch.abs(param_value), dim=(-1, -2))
+ scale = fp8_max / max_abs
+ scale_orig_shape = scale.shape
+ scale = scale.unsqueeze(-1).unsqueeze(-1)
+
+ # Quantize the weights
+ quantized_param = torch.clamp(param_value * scale, min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
+
+ quantized_param = quantized_param.permute(0, 1, 3, 2, 4)
+ # Reshape back to matrix shape
+ quantized_param = quantized_param.reshape(param_value_orig_shape)
+
+ # Reshape scale to match the number of blocks
+ scale = scale.reshape(scale_orig_shape).squeeze().reciprocal()
+
+ module._buffers[tensor_name] = quantized_param.to(target_device)
+ module._buffers["weight_scale_inv"] = scale.to(target_device)
+
+ def check_quantized_param(
+ self,
+ model: "PreTrainedModel",
+ param_value: "torch.Tensor",
+ param_name: str,
+ state_dict: Dict[str, Any],
+ **kwargs,
+ ):
+ from ..integrations.finegrained_fp8 import FP8Linear
+
+ module, tensor_name = get_module_from_name(model, param_name)
+
+ if isinstance(module, FP8Linear):
+ if self.pre_quantized or tensor_name == "bias":
+ if tensor_name == "weight" and param_value.dtype != torch.float8_e4m3fn:
+ raise ValueError("Expect quantized weights but got an unquantized weight")
+ return False
+ else:
+ if tensor_name == "weight_scale_inv":
+ raise ValueError("Expect unquantized weights but got a quantized weight_scale")
+ return True
+ return False
+
+ def _process_model_before_weight_loading(
+ self,
+ model: "PreTrainedModel",
+ keep_in_fp32_modules: Optional[List[str]] = None,
+ **kwargs,
+ ):
+ from ..integrations.finegrained_fp8 import replace_with_fp8_linear
+
+ self.modules_to_not_convert = self.get_modules_to_not_convert(
+ model, self.quantization_config.modules_to_not_convert, keep_in_fp32_modules
+ )
+
+ model = replace_with_fp8_linear(
+ model,
+ modules_to_not_convert=self.modules_to_not_convert,
+ quantization_config=self.quantization_config,
+ )
+
+ model.config.quantization_config = self.quantization_config
+
+ def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs):
+ return model
+
+ def update_missing_keys(self, model, missing_keys: List[str], prefix: str) -> List[str]:
+ from ..integrations import FP8Linear
+
+ not_missing_keys = []
+ for name, module in model.named_modules():
+ if isinstance(module, FP8Linear):
+ for missing in missing_keys:
+ if (
+ (name in missing or name in f"{prefix}.{missing}")
+ and not missing.endswith(".weight")
+ and not missing.endswith(".bias")
+ ):
+ not_missing_keys.append(missing)
+ return [k for k in missing_keys if k not in not_missing_keys]
+
+ def is_serializable(self, safe_serialization=None):
+ return True
+
+ @property
+ def is_trainable(self) -> bool:
+ return False
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_gptq.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_gptq.py
new file mode 100644
index 0000000000000000000000000000000000000000..0fc35bea1635edd957566f53005ff65a58d57282
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_gptq.py
@@ -0,0 +1,124 @@
+# Copyright 2024 The HuggingFace Inc. 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.
+import importlib
+from typing import TYPE_CHECKING, Optional
+
+from packaging import version
+
+from .base import HfQuantizer
+
+
+if TYPE_CHECKING:
+ from ..modeling_utils import PreTrainedModel
+
+from ..utils import is_auto_gptq_available, is_gptqmodel_available, is_optimum_available, is_torch_available, logging
+from ..utils.quantization_config import GPTQConfig, QuantizationConfigMixin
+
+
+if is_torch_available():
+ import torch
+
+logger = logging.get_logger(__name__)
+
+
+class GptqHfQuantizer(HfQuantizer):
+ """
+ Quantizer of the GPTQ method - for GPTQ the quantizer support calibration of the model through
+ `auto_gptq` or `gptqmodel` package. Quantization is done under the hood for users if they load a non-prequantized model.
+ """
+
+ requires_calibration = False
+ required_packages = ["optimum", "auto_gptq", "gptqmodel"]
+ optimum_quantizer = None
+
+ def __init__(self, quantization_config: QuantizationConfigMixin, **kwargs):
+ super().__init__(quantization_config, **kwargs)
+
+ if not is_optimum_available():
+ raise ImportError("Loading a GPTQ quantized model requires optimum (`pip install optimum`)")
+ from optimum.gptq import GPTQQuantizer
+
+ self.optimum_quantizer = GPTQQuantizer.from_dict(self.quantization_config.to_dict_optimum())
+
+ def validate_environment(self, *args, **kwargs):
+ if not is_optimum_available():
+ raise ImportError("Loading a GPTQ quantized model requires optimum (`pip install optimum`)")
+ if is_auto_gptq_available() and is_gptqmodel_available():
+ logger.warning("Detected gptqmodel and auto-gptq, will use gptqmodel")
+
+ gptq_supports_cpu = (
+ is_auto_gptq_available()
+ and version.parse(importlib.metadata.version("auto-gptq")) > version.parse("0.4.2")
+ ) or is_gptqmodel_available()
+ if not gptq_supports_cpu and not torch.cuda.is_available():
+ raise RuntimeError("GPU is required to quantize or run quantize model.")
+ elif not (is_auto_gptq_available() or is_gptqmodel_available()):
+ raise ImportError(
+ "Loading a GPTQ quantized model requires gptqmodel (`pip install gptqmodel`) or auto-gptq (`pip install auto-gptq`) library. "
+ )
+ elif is_auto_gptq_available() and version.parse(importlib.metadata.version("auto_gptq")) < version.parse(
+ "0.4.2"
+ ):
+ raise ImportError(
+ "You need a version of auto_gptq >= 0.4.2 to use GPTQ: `pip install --upgrade auto-gptq` or use gptqmodel by `pip install gptqmodel>=1.4.3`."
+ )
+ elif is_gptqmodel_available() and (
+ version.parse(importlib.metadata.version("gptqmodel")) < version.parse("1.4.3")
+ or version.parse(importlib.metadata.version("optimum")) < version.parse("1.23.99")
+ ):
+ raise ImportError("The gptqmodel version should be >= 1.4.3, optimum version should >= 1.24.0")
+
+ def update_torch_dtype(self, torch_dtype: "torch.dtype") -> "torch.dtype":
+ if torch_dtype is None:
+ torch_dtype = torch.float16
+ logger.info("Loading the model in `torch.float16`. To overwrite it, set `torch_dtype` manually.")
+ elif torch_dtype != torch.float16:
+ logger.info("We suggest you to set `torch_dtype=torch.float16` for better efficiency with GPTQ.")
+ return torch_dtype
+
+ def update_device_map(self, device_map):
+ if device_map is None:
+ device_map = {"": torch.device("cpu")}
+ # Only with auto-gptq do not support CPU, we should move the model to cuda if available.
+ if not is_gptqmodel_available() and device_map in ("cpu", {"": torch.device("cpu")}):
+ device_map == {"": 0}
+ return device_map
+
+ def _process_model_before_weight_loading(self, model: "PreTrainedModel", **kwargs):
+ if model.__class__.main_input_name != "input_ids":
+ raise RuntimeError("We can only quantize pure text model.")
+
+ if self.pre_quantized:
+ # compat: latest optimum has gptqmodel refactor
+ if version.parse(importlib.metadata.version("optimum")) <= version.parse("1.23.99"):
+ model = self.optimum_quantizer.convert_model(model)
+ else:
+ model = self.optimum_quantizer.convert_model(model, **kwargs)
+
+ def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs):
+ if self.pre_quantized:
+ model = self.optimum_quantizer.post_init_model(model)
+ else:
+ if self.quantization_config.tokenizer is None:
+ self.quantization_config.tokenizer = model.name_or_path
+
+ self.optimum_quantizer.quantize_model(model, self.quantization_config.tokenizer)
+ model.config.quantization_config = GPTQConfig.from_dict(self.optimum_quantizer.to_dict())
+
+ @property
+ def is_trainable(self, model: Optional["PreTrainedModel"] = None):
+ return True
+
+ def is_serializable(self, safe_serialization=None):
+ return True
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_higgs.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_higgs.py
new file mode 100644
index 0000000000000000000000000000000000000000..83c102f16c3e303514d856508e4dc37d87163ffb
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_higgs.py
@@ -0,0 +1,202 @@
+# Copyright 2024 The HuggingFace Inc. 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.
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
+
+from ..utils.logging import tqdm
+from .base import HfQuantizer
+from .quantizers_utils import get_module_from_name
+
+
+if TYPE_CHECKING:
+ from ..modeling_utils import PreTrainedModel
+
+from ..utils import is_accelerate_available, is_flute_available, is_hadamard_available, is_torch_available, logging
+from ..utils.quantization_config import QuantizationConfigMixin
+
+
+if is_torch_available():
+ import torch
+
+logger = logging.get_logger(__name__)
+
+
+class HiggsHfQuantizer(HfQuantizer):
+ """
+ Quantizer of the HIGGS method. Enables the loading of prequantized models and in-flight quantization of full-precision models.
+ """
+
+ requires_calibration = False
+ requires_parameters_quantization = True
+ required_packages = ["flute-kernel", "fast_hadamard_transform"]
+
+ def __init__(self, quantization_config: QuantizationConfigMixin, **kwargs):
+ super().__init__(quantization_config, **kwargs)
+ self.quantization_config = quantization_config
+
+ def validate_environment(self, device_map, **kwargs):
+ if not torch.cuda.is_available():
+ raise NotImplementedError("HIGGS quantization is only supported on GPU. Please use a different quantizer.")
+
+ if not is_accelerate_available():
+ raise ImportError("Using `higgs` quantization requires Accelerate: `pip install accelerate`")
+
+ if not is_flute_available():
+ raise ImportError("Using `higgs` quantization requires FLUTE: `pip install flute-kernel>=0.3.0`")
+
+ if not is_hadamard_available():
+ raise ImportError(
+ "Using `higgs` quantization requires fast_hadamard_transform: `pip install fast_hadamard_transform`"
+ )
+
+ if device_map is None:
+ raise ValueError(
+ "You are attempting to load a HIGGS model without setting device_map."
+ " Please set device_map comprised of 'cuda' devices."
+ )
+ elif isinstance(device_map, dict) and ("cpu" in device_map.values() or "disk" in device_map.values()):
+ raise ValueError(
+ "You are attempting to load a HIGGS model with a device_map that contains a CPU or disk device."
+ " This is not supported. Please remove the CPU or disk device from the device_map."
+ )
+
+ def update_torch_dtype(self, torch_dtype: "torch.dtype") -> "torch.dtype":
+ if torch_dtype is None:
+ logger.info("`torch_dtype` is None. Setting `torch_dtype=torch.float16` for FLUTE compatibility.")
+ torch_dtype = torch.float16
+ elif torch_dtype != torch.float16 and torch_dtype != torch.bfloat16:
+ raise ValueError(
+ f"Invalid `torch_dtype` {torch_dtype}. HIGGS quantization only supports `torch_dtype=torch.float16` or `torch_dtype=torch.bfloat16`."
+ )
+
+ return torch_dtype
+
+ def create_quantized_param(
+ self,
+ model: "PreTrainedModel",
+ param_value: "torch.Tensor",
+ param_name: str,
+ target_device: "torch.device",
+ state_dict: Dict[str, Any],
+ unexpected_keys: Optional[List[str]] = None,
+ ):
+ from ..integrations import quantize_with_higgs
+
+ """
+ Quantizes weights into weight and weight_scale
+ """
+ flute_dict = quantize_with_higgs(
+ param_value.to(target_device),
+ self.quantization_config.bits,
+ self.quantization_config.p,
+ self.quantization_config.group_size,
+ self.quantization_config.hadamard_size,
+ )
+ del param_value
+
+ module, _ = get_module_from_name(model, param_name)
+ module_name = ".".join(param_name.split(".")[:-1])
+ for key, value in flute_dict.items():
+ if key in module._parameters:
+ module._parameters[key] = torch.nn.Parameter(value, requires_grad=False)
+ elif key in module._buffers:
+ module._buffers[key] = torch.nn.Buffer(value)
+ elif key == "tune_metadata":
+ module.tune_metadata = value
+ self.quantization_config.tune_metadata[module_name] = value.to_dict()
+ else:
+ raise ValueError(f"Unexpected key {key} in module {module}")
+
+ if unexpected_keys is not None and param_name in unexpected_keys:
+ unexpected_keys.remove(param_name)
+
+ def _process_model_before_weight_loading(
+ self,
+ model: "PreTrainedModel",
+ **kwargs,
+ ):
+ from ..integrations import replace_with_higgs_linear
+
+ replace_with_higgs_linear(
+ model,
+ quantization_config=self.quantization_config,
+ )
+ model.config.quantization_config = self.quantization_config
+
+ def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs):
+ from flute.tune import TuneMetaData, maybe_tune_and_repack
+ from flute.utils import make_workspace_streamk
+
+ from ..integrations import HiggsLinear
+
+ flute_workspaces = {}
+ flute_modules = {name: module for name, module in model.named_modules() if isinstance(module, HiggsLinear)}
+ for name, module in tqdm(flute_modules.items(), desc="Repacking HIGGS modules", leave=False):
+ # Every HiggsLinear needs a "workspace": a buffer for the unpacking operation.
+ # This buffer needs to be on the same device as the weights, but can be reused across modules otherwise.
+ if module.weight.device not in flute_workspaces:
+ flute_workspaces[module.weight.device] = make_workspace_streamk(device=module.weight.device)
+ module.workspace = flute_workspaces[module.weight.device]
+
+ # FLUTE weights are packed in a way that is optimized for a specific number of SMs (GPU streaming multiprocessors).
+ # If the model is loaded on a different device than the one it was saved on, we need to repack the weights.
+ module.tune_metadata = TuneMetaData.from_dict(self.quantization_config.tune_metadata[name])
+ module.weight.data, module.tune_metadata = maybe_tune_and_repack(
+ weight=module.weight.data,
+ scales=module.scales.data,
+ metadata=module.tune_metadata,
+ )
+ self.quantization_config.tune_metadata[name] = module.tune_metadata.to_dict()
+
+ def update_missing_keys(self, model, missing_keys: List[str], prefix: str) -> List[str]:
+ from ..integrations import HiggsLinear
+
+ higgs_names = {name for name, module in model.named_modules() if isinstance(module, HiggsLinear)}
+
+ def should_update(key: str) -> bool:
+ if key.endswith(".weight") or key.endswith(".bias"):
+ return False
+ full_key = f"{prefix}.{key}"
+ return any(name in key or name in full_key for name in higgs_names)
+
+ return [key for key in missing_keys if not should_update(key)]
+
+ @property
+ def is_trainable(self, model: Optional["PreTrainedModel"] = None):
+ return False
+
+ def is_serializable(self, safe_serialization=None):
+ return True
+
+ def check_quantized_param(
+ self,
+ model: "PreTrainedModel",
+ param_value: "torch.Tensor",
+ param_name: str,
+ state_dict: Dict[str, Any],
+ **kwargs,
+ ) -> bool:
+ from ..integrations import HiggsLinear
+
+ module, tensor_name = get_module_from_name(model, param_name)
+ if isinstance(module, HiggsLinear) and tensor_name == "weight" and param_value.dtype != torch.int16:
+ # Only quantize weights of HiggsLinear modules that are not already quantized
+ return True
+ else:
+ return False
+
+ def _dequantize(self, model):
+ from ..integrations import dequantize_higgs
+
+ model = dequantize_higgs(model)
+ return model
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_hqq.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_hqq.py
new file mode 100644
index 0000000000000000000000000000000000000000..4adc323f958ebfcd7de8fd5d62d8fdcb99606582
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_hqq.py
@@ -0,0 +1,322 @@
+# Copyright 2024 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.
+
+from typing import TYPE_CHECKING, Any, Dict, List
+
+from ..integrations import prepare_for_hqq_linear
+from ..utils import is_accelerate_available, is_hqq_available, is_torch_available, logging
+from .base import HfQuantizer
+from .quantizers_utils import get_module_from_name
+
+
+if TYPE_CHECKING:
+ from ..modeling_utils import PreTrainedModel
+
+
+if is_accelerate_available():
+ from accelerate.hooks import remove_hook_from_module
+
+if is_torch_available():
+ import torch
+
+logger = logging.get_logger(__name__)
+
+
+# Finds the parent of a node module named "name"
+def find_parent(model, name):
+ module_tree = name.split(".")[:-1]
+ parent = model
+ for m in module_tree:
+ parent = parent._modules[m]
+ return parent
+
+
+class HqqHfQuantizer(HfQuantizer):
+ """
+ HQQ quantizer base HF class.
+ nn.Linear modules are first tagged with quant_config in _process_model_before_weight_loading().
+ The actual quantization and offloading to the GPU is done in check_quantized_param().
+ """
+
+ use_keep_in_fp32_modules = False
+ requires_parameters_quantization = True
+ requires_calibration = False
+ required_packages = ["hqq"]
+
+ def __init__(self, quantization_config, **kwargs):
+ super().__init__(quantization_config, **kwargs)
+ self.torch_dtype = None
+ self.using_multi_gpu = False
+
+ def validate_environment(self, *args, **kwargs):
+ if not (is_hqq_available()):
+ raise ImportError(
+ "A valid HQQ version (>=0.2.1) is not available. Please follow the instructions to install it: `https://github.com/mobiusml/hqq/`."
+ )
+
+ if kwargs.get("from_tf", False) or kwargs.get("from_flax", False):
+ raise ValueError(
+ "Converting weights from tf/flax weights is currently not supported, please make"
+ " sure the weights are in PyTorch format."
+ )
+
+ if not torch.cuda.is_available():
+ raise RuntimeError("No GPU found. A GPU is needed for quantization.")
+
+ if self.torch_dtype is None:
+ if "torch_dtype" in kwargs:
+ self.torch_dtype = kwargs["torch_dtype"]
+ else:
+ self.torch_dtype = torch.float32
+ logger.info("Setting torch_dtype to torch.float32 as the default value since it was not specified.")
+
+ device_map = kwargs.get("device_map", None)
+ if isinstance(device_map, dict):
+ if "cpu" in device_map.values() or "disk" in device_map.values():
+ raise ValueError(
+ "You are attempting to use an HQQ model with a device_map that contains a CPU or disk device."
+ " This is not supported. Please remove the CPU or disk device from the device_map."
+ )
+ else:
+ self.using_multi_gpu = len(set(device_map.values())) > 1
+
+ def update_missing_keys(
+ self, model: "PreTrainedModel", missing_keys: List[str], prefix: str, **kwargs
+ ) -> List[str]:
+ if self.pre_quantized:
+ return [key for key in missing_keys if ("weight" not in key)]
+ else:
+ return missing_keys
+
+ # Adds missing keys for HQQLinear modules that are loaded but the model with initialized with torch.nn.Linear
+ def update_expected_keys(
+ self, model: "PreTrainedModel", expected_keys: List[str], loaded_keys: List[str]
+ ) -> List[str]:
+ if not self.pre_quantized:
+ return expected_keys
+
+ # Collects all quantizable (linear) layers
+ def _find_hqq_quantizable_layers(model, layers):
+ for name, module in model.named_children():
+ if isinstance(module, (torch.nn.Linear)):
+ layers.add(module.name)
+ _find_hqq_quantizable_layers(module, layers)
+
+ new_keys = set(expected_keys)
+ if is_hqq_available():
+ from hqq.core.quantize import HQQLinear
+
+ # Name modules
+ for name, module in model.named_modules():
+ module.name = name
+
+ # valid modules are Linear layers that have HQQLinear state_dict. We ignore skip_modules and any layers with Linear state_dict() params
+ _valid_modules = set()
+ _find_hqq_quantizable_layers(model, _valid_modules)
+
+ # Remove skipped modules
+ _skipped_modules = set()
+ for _module in _valid_modules:
+ for _skip_module in model.config.quantization_config["skip_modules"]:
+ if _skip_module in _module:
+ _skipped_modules.add(_module)
+ _valid_modules -= _skipped_modules
+
+ # Append new expected layers based on _ref_keys
+ _ref_keys = HQQLinear(
+ linear_layer=None, quant_config=None, compute_dtype=torch.float16, device="cpu"
+ ).state_dict_keys() - {"bias"}
+
+ # Clean-up
+ _rm_keys = set()
+ for key in new_keys:
+ if any(_module in key for _module in _valid_modules):
+ _rm_keys.add(key)
+ new_keys -= _rm_keys
+ # At this point, new_keys contains all the keys of the layers that are NOT HQQLinear or torch.nn.Linear
+
+ # Re-populate Linear/HQQLinear
+ for _module in _valid_modules:
+ if _module + ".weight" in loaded_keys:
+ new_keys.add(_module + ".weight")
+ else:
+ new_keys.update({_module + "." + _ref_key for _ref_key in _ref_keys})
+ if _module + ".bias" in loaded_keys:
+ new_keys.add(_module + ".bias")
+
+ return list(new_keys)
+
+ def check_quantized_param(
+ self,
+ model: "PreTrainedModel",
+ param_value: "torch.Tensor",
+ param_name: str,
+ state_dict: Dict[str, Any],
+ **kwargs,
+ ) -> bool:
+ if is_hqq_available():
+ from hqq.core.quantize import HQQLinear
+ module, tensor_name = get_module_from_name(model, param_name)
+
+ if self.pre_quantized:
+ return (
+ (isinstance(module, torch.nn.Linear) or isinstance(module, HQQLinear))
+ and tensor_name != "weight"
+ and tensor_name != "bias"
+ )
+ else:
+ # we need a special path for bias since hqq overwrote load_state_dict for this layer
+ return (
+ isinstance(module, torch.nn.Linear)
+ and tensor_name == "weight"
+ or (isinstance(module, HQQLinear) and tensor_name == "bias")
+ )
+
+ def create_quantized_param(
+ self,
+ model: "PreTrainedModel",
+ param_value: "torch.Tensor",
+ param_name: str,
+ target_device: "torch.device",
+ state_dict: Dict[str, Any],
+ unexpected_keys: List[str],
+ ):
+ """
+ Each nn.Linear layer is processed here.
+ We first check if the corresponding module state_dict contains already HQQ quantized parameters.
+ If not, we create a temp linear layer with the module state_dict params and use it for quantization
+ """
+
+ if is_hqq_available():
+ from hqq.core.quantize import HQQLinear
+
+ module, tensor_name = get_module_from_name(model, param_name)
+ layer_name = ".".join(param_name.split(".")[:-1])
+ parent_module = find_parent(model, layer_name)
+ node = layer_name.split(".")[-1]
+
+ if tensor_name == "bias":
+ # this should already be set
+ return
+
+ # set module state_dict
+ module_state_dict = {}
+ for k, v in state_dict.items():
+ if layer_name + "." in k:
+ module_state_dict[k.split(".")[-1]] = v
+ if unexpected_keys is not None and k in unexpected_keys:
+ unexpected_keys.remove(k)
+
+ if self.pre_quantized:
+ if isinstance(module, HQQLinear):
+ return
+ else:
+ hqq_layer = HQQLinear(
+ linear_layer=None,
+ quant_config=None,
+ compute_dtype=self.torch_dtype,
+ device=target_device,
+ )
+
+ hqq_layer.load_state_dict(module_state_dict)
+
+ if hqq_layer.bias is not None and isinstance(hqq_layer.bias, torch.Tensor):
+ hqq_layer.bias = torch.nn.Parameter(hqq_layer.bias)
+
+ if self.using_multi_gpu:
+ hqq_layer = self._patch_layer_for_multigpu(hqq_layer)
+
+ setattr(parent_module, node, hqq_layer)
+
+ # cleanup
+ del module.__dict__, module
+ torch.cuda.empty_cache()
+ return
+
+ # Step 1: populate module with weight/bias from module state dict
+ for key in module_state_dict:
+ setattr(module, key, torch.nn.Parameter(module_state_dict[key]))
+
+ # Step 2: Replace module with either HQQLinear or move it to device. We do this via setattr on the parent as doing on it on the module
+ # directly doesn't work.
+ quant_config = model.config.quantization_config["quant_config"]
+ skip_modules = model.config.quantization_config["skip_modules"]
+ module_tag = ".".join(module.name.split(".")[-2:])
+ module_quant_config = None
+ if "weight_quant_params" in quant_config:
+ module_quant_config = quant_config
+ elif module_tag in quant_config:
+ module_quant_config = quant_config[module_tag]
+
+ for skip_module in skip_modules:
+ if skip_module in module.name:
+ module_quant_config = None
+ break
+
+ if module_quant_config is not None:
+ hqq_layer = HQQLinear(
+ module,
+ quant_config=module_quant_config,
+ compute_dtype=self.torch_dtype,
+ device=target_device,
+ del_orig=True,
+ )
+
+ if hqq_layer.bias is not None and isinstance(hqq_layer.bias, torch.Tensor):
+ hqq_layer.bias = torch.nn.Parameter(hqq_layer.bias)
+
+ if self.using_multi_gpu:
+ hqq_layer = self._patch_layer_for_multigpu(hqq_layer)
+
+ setattr(parent_module, node, hqq_layer)
+
+ else:
+ module = module.to(dtype=self.torch_dtype, device=target_device)
+ setattr(parent_module, node, module)
+
+ torch.cuda.empty_cache()
+
+ # Remove accelerate hook and uses a simpler forward pass. Otherwise, this breaks with multi-gpu
+ def _patch_layer_for_multigpu(self, hqq_layer):
+ hqq_layer = remove_hook_from_module(hqq_layer)
+
+ def forward_with_device(self, x):
+ out = torch.matmul(x.to(self.device), self.dequantize().t())
+ if self.bias is not None:
+ out += self.bias
+ return out
+
+ hqq_layer.forward = lambda x: forward_with_device(hqq_layer, x)
+ return hqq_layer
+
+ def _process_model_before_weight_loading(
+ self,
+ model: "PreTrainedModel",
+ **kwargs,
+ ):
+ # Add the corresponding quant_config to each valid module. This allows us to do the actual nn.Linear -> HQQLinear conversion in create_quantized_param().
+ # prepare_for_hqq_linear() also sets the right quantization config inside the model (model.config.quantization_config) and the layers (hqq_layer.quant_config)
+ model = prepare_for_hqq_linear(model, quantization_config=self.quantization_config)
+
+ def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs):
+ model.is_hqq_quantized = True
+ model.is_hqq_serializable = self.is_serializable()
+ return model
+
+ def is_serializable(self, safe_serialization=None):
+ return True
+
+ @property
+ def is_trainable(self) -> bool:
+ return True
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_quanto.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_quanto.py
new file mode 100644
index 0000000000000000000000000000000000000000..be760f0d430c16c5a81036972689ebca553bb20f
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_quanto.py
@@ -0,0 +1,201 @@
+# Copyright 2024 The HuggingFace Inc. 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.
+import importlib
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+
+from packaging import version
+
+from .base import HfQuantizer
+from .quantizers_utils import get_module_from_name
+
+
+if TYPE_CHECKING:
+ from ..modeling_utils import PreTrainedModel
+
+from ..utils import (
+ is_accelerate_available,
+ is_optimum_quanto_available,
+ is_torch_available,
+ logging,
+)
+from ..utils.quantization_config import QuantoConfig
+
+
+if is_torch_available():
+ import torch
+
+logger = logging.get_logger(__name__)
+
+
+class QuantoHfQuantizer(HfQuantizer):
+ """
+ Quantizer for the quanto library
+ """
+
+ required_packages = ["quanto", "accelerate"]
+ requires_parameters_quantization = True
+ requires_calibration = False
+
+ def __init__(self, quantization_config: QuantoConfig, **kwargs):
+ super().__init__(quantization_config, **kwargs)
+ self.post_init()
+
+ def post_init(self):
+ r"""
+ Safety checker
+ """
+ if self.quantization_config.activations is not None and not self.pre_quantized:
+ raise ValueError(
+ "We don't support quantizing the activations with transformers library."
+ "Use quanto library for more complex use cases such as activations quantization, calibration and quantization aware training."
+ )
+
+ def validate_environment(self, *args, **kwargs):
+ if not is_optimum_quanto_available():
+ raise ImportError(
+ "Loading an optimum-quanto quantized model requires optimum-quanto library (`pip install optimum-quanto`)"
+ )
+ if not is_accelerate_available():
+ raise ImportError(
+ "Loading an optimum-quanto quantized model requires accelerate library (`pip install accelerate`)"
+ )
+
+ def update_device_map(self, device_map):
+ if device_map is None:
+ device_map = {"": "cpu"}
+ logger.info(
+ "The device_map was not initialized. "
+ "Setting device_map to {'':'cpu'}. "
+ "If you want to use the model for inference, please set device_map ='auto'"
+ )
+ return device_map
+
+ def update_torch_dtype(self, torch_dtype: "torch.dtype") -> "torch.dtype":
+ if torch_dtype is None:
+ logger.info("You did not specify `torch_dtype` in `from_pretrained`. Setting it to `torch.float32`.")
+ torch_dtype = torch.float32
+ return torch_dtype
+
+ def update_missing_keys(self, model, missing_keys: List[str], prefix: str) -> List[str]:
+ if is_optimum_quanto_available():
+ from optimum.quanto import QModuleMixin
+
+ not_missing_keys = []
+ for name, module in model.named_modules():
+ if isinstance(module, QModuleMixin):
+ for missing in missing_keys:
+ if (
+ (name in missing or name in f"{prefix}.{missing}")
+ and not missing.endswith(".weight")
+ and not missing.endswith(".bias")
+ ):
+ not_missing_keys.append(missing)
+ return [k for k in missing_keys if k not in not_missing_keys]
+
+ def check_quantized_param(
+ self,
+ model: "PreTrainedModel",
+ param_value: "torch.Tensor",
+ param_name: str,
+ state_dict: Dict[str, Any],
+ **kwargs,
+ ) -> bool:
+ """
+ Check if a parameter needs to be quantized.
+ """
+ if is_optimum_quanto_available():
+ from optimum.quanto import QModuleMixin
+
+ device_map = kwargs.get("device_map", None)
+ param_device = kwargs.get("param_device", None)
+ # we don't quantize the model if the module is going to be offloaded to the cpu
+ if device_map is not None and param_device is not None:
+ device_map_values = set(device_map.values())
+ if param_device == "cpu" and len(device_map_values) > 1:
+ if not (device_map_values == {"cpu"} or device_map_values == {"cpu", "disk"}):
+ return False
+
+ module, tensor_name = get_module_from_name(model, param_name)
+ # We only quantize the weights and the bias is not quantized.
+ if isinstance(module, QModuleMixin) and "weight" in tensor_name:
+ # if the weights are quantized, don't need to recreate it again with `create_quantized_param`
+ return not module.frozen
+ else:
+ return False
+
+ def adjust_max_memory(self, max_memory: Dict[str, Union[int, str]]) -> Dict[str, Union[int, str]]:
+ max_memory = {key: val * 0.90 for key, val in max_memory.items()}
+ return max_memory
+
+ def create_quantized_param(
+ self,
+ model: "PreTrainedModel",
+ param_value: "torch.Tensor",
+ param_name: str,
+ target_device: "torch.device",
+ *args,
+ **kwargs,
+ ):
+ """
+ Create the quantized parameter by calling .freeze() after setting it to the module.
+ """
+ from accelerate.utils import set_module_tensor_to_device
+
+ set_module_tensor_to_device(model, param_name, target_device, param_value)
+ module, _ = get_module_from_name(model, param_name)
+ module.freeze()
+ module.weight.requires_grad = False
+
+ def adjust_target_dtype(self, target_dtype: "torch.dtype") -> "torch.dtype":
+ if version.parse(importlib.metadata.version("accelerate")) > version.parse("0.27.0"):
+ from accelerate.utils import CustomDtype
+
+ mapping = {
+ "int8": torch.int8,
+ "float8": CustomDtype.FP8,
+ "int4": CustomDtype.INT4,
+ "int2": CustomDtype.INT2,
+ }
+ target_dtype = mapping[self.quantization_config.weights]
+ return target_dtype
+ else:
+ raise ValueError(
+ "You are using `device_map='auto'` on an optimum-quanto quantized model. To automatically compute"
+ " the appropriate device map, you should upgrade your `accelerate` library,"
+ "`pip install --upgrade accelerate` or install it from source."
+ )
+
+ def _process_model_before_weight_loading(
+ self, model: "PreTrainedModel", keep_in_fp32_modules: Optional[List[str]] = None, **kwargs
+ ):
+ from ..integrations import replace_with_quanto_layers
+
+ self.modules_to_not_convert = self.get_modules_to_not_convert(
+ model, self.quantization_config.modules_to_not_convert, keep_in_fp32_modules
+ )
+
+ model, _ = replace_with_quanto_layers(
+ model, modules_to_not_convert=self.modules_to_not_convert, quantization_config=self.quantization_config
+ )
+ model.config.quantization_config = self.quantization_config
+
+ def _process_model_after_weight_loading(self, model, **kwargs):
+ return model
+
+ @property
+ def is_trainable(self, model: Optional["PreTrainedModel"] = None):
+ return True
+
+ def is_serializable(self, safe_serialization=None):
+ return False
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_quark.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_quark.py
new file mode 100644
index 0000000000000000000000000000000000000000..374360b1cb8f20724d74fe3d59cb78b0b47b2912
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_quark.py
@@ -0,0 +1,113 @@
+# coding=utf-8
+# Copyright 2025 Advanced Micro Devices, Inc. and The HuggingFace Inc. 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.
+
+from typing import TYPE_CHECKING, Any, Dict
+
+from ..file_utils import is_torch_available
+from .base import HfQuantizer
+
+
+if TYPE_CHECKING:
+ from ..modeling_utils import PreTrainedModel
+
+ if is_torch_available():
+ import torch
+
+from ..utils import is_accelerate_available, is_quark_available, logging
+
+
+if is_accelerate_available():
+ from accelerate.utils import set_module_tensor_to_device
+
+logger = logging.get_logger(__name__)
+
+
+CHECKPOINT_KEYS = {
+ "weight_scale": "weight_quantizer.scale",
+ "bias_scale": "bias_quantizer.scale",
+ "input_scale": "input_quantizer.scale",
+ "output_scale": "output_quantizer.scale",
+ "weight_zero_point": "weight_quantizer.zero_point",
+ "bias_zero_point": "bias_quantizer.zero_point",
+ "input_zero_point": "input_quantizer.zero_point",
+ "output_zero_point": "output_quantizer.zero_point",
+}
+
+
+class QuarkHfQuantizer(HfQuantizer):
+ """
+ Quark quantizer (https://quark.docs.amd.com/latest/).
+ """
+
+ requires_calibration = True # On-the-fly quantization with quark is not supported for now.
+ required_packages = ["quark"]
+
+ # Checkpoints are expected to be already quantized when loading a quark model. However, as some keys from
+ # the checkpoint might mismatch the model parameters keys, we use the `create_quantized_param` method
+ # to load the checkpoints, remapping the keys.
+ requires_parameters_quantization = True
+
+ def __init__(self, quantization_config, **kwargs):
+ super().__init__(quantization_config, **kwargs)
+
+ self.json_export_config = quantization_config.json_export_config
+
+ def validate_environment(self, *args, **kwargs):
+ if not is_quark_available():
+ raise ImportError(
+ "Loading a Quark quantized model requires the `quark` library but it was not found in the environment. Please refer to https://quark.docs.amd.com/latest/install.html."
+ )
+
+ def _process_model_before_weight_loading(self, model: "PreTrainedModel", **kwargs):
+ from quark.torch.export.api import _map_to_quark
+
+ _map_to_quark(
+ model,
+ self.quantization_config.quant_config,
+ pack_method=self.json_export_config.pack_method,
+ custom_mode=self.quantization_config.custom_mode,
+ )
+
+ return model
+
+ def check_quantized_param(
+ self,
+ model: "PreTrainedModel",
+ param_value: "torch.Tensor",
+ param_name: str,
+ state_dict: Dict[str, Any],
+ **kwargs,
+ ) -> bool:
+ return True
+
+ def create_quantized_param(
+ self, model, param, param_name, param_device, state_dict, unexpected_keys
+ ) -> "torch.nn.Parameter":
+ postfix = param_name.split(".")[-1]
+
+ if postfix in CHECKPOINT_KEYS:
+ param_name = param_name.replace(postfix, CHECKPOINT_KEYS[postfix])
+
+ set_module_tensor_to_device(model, param_name, param_device, value=param)
+
+ def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs):
+ return model
+
+ def is_serializable(self, safe_serialization=None):
+ return False
+
+ @property
+ def is_trainable(self):
+ return False
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_spqr.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_spqr.py
new file mode 100644
index 0000000000000000000000000000000000000000..4cf1193edbf2cc676704db06fda7505c824548c0
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_spqr.py
@@ -0,0 +1,90 @@
+# Copyright 2024 The HuggingFace Inc. 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/lic enses/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.
+from typing import TYPE_CHECKING, List, Optional
+
+from .base import HfQuantizer
+
+
+if TYPE_CHECKING:
+ from ..modeling_utils import PreTrainedModel
+
+from ..integrations import replace_with_spqr_linear
+from ..utils import is_accelerate_available, is_spqr_available, is_torch_available, logging
+from ..utils.quantization_config import QuantizationConfigMixin
+
+
+if is_torch_available():
+ import torch
+
+logger = logging.get_logger(__name__)
+
+
+class SpQRHfQuantizer(HfQuantizer):
+ """
+ Quantizer of the SpQR method. Enables the loading of prequantized models.
+ """
+
+ requires_calibration = True
+
+ def __init__(self, quantization_config: QuantizationConfigMixin, **kwargs):
+ super().__init__(quantization_config, **kwargs)
+ self.quantization_config = quantization_config
+
+ def validate_environment(self, *args, **kwargs):
+ if not torch.cuda.is_available():
+ raise RuntimeError("GPU is required to run SpQR quantized model.")
+
+ if not is_accelerate_available():
+ raise ImportError("Using `spqr` quantization requires Accelerate: `pip install accelerate`")
+
+ if not is_spqr_available():
+ raise ImportError("Using `spqr` quantization requires SpQR: `pip install spqr_quant[gpu]`")
+
+ def update_torch_dtype(self, torch_dtype: "torch.dtype") -> "torch.dtype":
+ if torch_dtype is None:
+ torch_dtype = torch.float16
+ logger.info("Assuming SpQR inference on GPU and loading the model in `torch.float16`.")
+ elif torch_dtype != torch.float16:
+ raise ValueError(
+ "You cannot use any type other than torch.float16 for SpQR. Please either leave it None or set it to"
+ "torch.float16 explicitly."
+ )
+ return torch_dtype
+
+ def _process_model_before_weight_loading(
+ self,
+ model: "PreTrainedModel",
+ keep_in_fp32_modules: Optional[List[str]] = None,
+ **kwargs,
+ ):
+ self.modules_to_not_convert = self.get_modules_to_not_convert(
+ model, self.quantization_config.modules_to_not_convert, keep_in_fp32_modules
+ )
+
+ replace_with_spqr_linear(
+ model,
+ quantization_config=self.quantization_config,
+ modules_to_not_convert=self.modules_to_not_convert,
+ )
+ model.config.quantization_config = self.quantization_config
+
+ def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs):
+ return model
+
+ @property
+ def is_trainable(self, model: Optional["PreTrainedModel"] = None):
+ return False
+
+ def is_serializable(self, safe_serialization=None):
+ return True
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_torchao.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_torchao.py
new file mode 100644
index 0000000000000000000000000000000000000000..ac8155f22767dc77d06343fe661b675cc02eb2da
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_torchao.py
@@ -0,0 +1,286 @@
+# Copyright 2024 The HuggingFace Inc. 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.
+import importlib
+import re
+import types
+from typing import TYPE_CHECKING, Optional, Union
+
+from packaging import version
+
+from .base import HfQuantizer
+from .quantizers_utils import get_module_from_name
+
+
+if TYPE_CHECKING:
+ from ..modeling_utils import PreTrainedModel
+
+from typing import Any, Dict, List
+
+from ..utils import is_torch_available, is_torchao_available, logging
+from ..utils.quantization_config import TorchAoConfig
+
+
+if is_torch_available():
+ import torch
+ import torch.nn as nn
+
+logger = logging.get_logger(__name__)
+
+
+def fuzzy_match_size(config_name: str) -> Optional[str]:
+ """
+ Extract the size digit from strings like "4weight", "8weight".
+ Returns the digit as an integer if found, otherwise None.
+ """
+ config_name = config_name.lower()
+
+ str_match = re.search(r"(\d)weight", config_name)
+
+ if str_match:
+ return str_match.group(1)
+
+ return None
+
+
+# Finds the parent of a node module named "name"
+def find_parent(model, name):
+ module_tree = name.split(".")[:-1]
+ parent = model
+ for m in module_tree:
+ parent = parent._modules[m]
+ return parent
+
+
+def _quantization_type(weight):
+ from torchao.dtypes import AffineQuantizedTensor
+ from torchao.quantization.linear_activation_quantized_tensor import LinearActivationQuantizedTensor
+
+ if isinstance(weight, AffineQuantizedTensor):
+ return f"{weight.__class__.__name__}({weight._quantization_type()})"
+
+ if isinstance(weight, LinearActivationQuantizedTensor):
+ return f"{weight.__class__.__name__}(activation={weight.input_quant_func}, weight={_quantization_type(weight.original_weight_tensor)})"
+
+
+def _linear_extra_repr(self):
+ weight = _quantization_type(self.weight)
+ if weight is None:
+ return f"in_features={self.weight.shape[1]}, out_features={self.weight.shape[0]}, weight=None"
+ else:
+ return f"in_features={self.weight.shape[1]}, out_features={self.weight.shape[0]}, weight={weight}"
+
+
+class TorchAoHfQuantizer(HfQuantizer):
+ """
+ Quantizer for torchao: https://github.com/pytorch/ao/
+ """
+
+ requires_parameters_quantization = True
+ requires_calibration = False
+ required_packages = ["torchao"]
+
+ def __init__(self, quantization_config, **kwargs):
+ super().__init__(quantization_config, **kwargs)
+
+ def validate_environment(self, *args, **kwargs):
+ if not is_torchao_available():
+ raise ImportError("Loading an torchao quantized model requires torchao library (`pip install torchao`)")
+
+ self.offload = False
+ device_map = kwargs.get("device_map", None)
+ if isinstance(device_map, dict):
+ if "cpu" in device_map.values() or "disk" in device_map.values():
+ if self.pre_quantized:
+ raise ValueError(
+ "You are attempting to perform cpu/disk offload with a pre-quantized torchao model "
+ "This is not supported yet . Please remove the CPU or disk device from the device_map."
+ )
+ else:
+ self.offload = True
+ if self.pre_quantized:
+ weights_only = kwargs.get("weights_only", None)
+ if weights_only:
+ torch_version = version.parse(importlib.metadata.version("torch"))
+ if torch_version < version.parse("2.5.0"):
+ raise RuntimeError(
+ f"In order to use torchao pre-quantized model, you need to have torch>=2.5.0. However, the current version is {torch_version}."
+ f" You can also set with `weights_only=False` in `from_pretrained` if you don't want to update torch"
+ )
+
+ def update_torch_dtype(self, torch_dtype):
+ if self.quantization_config.quant_type == "int4_weight_only":
+ if torch_dtype is not None and torch_dtype != torch.bfloat16:
+ logger.warning_once(
+ f"Setting torch_dtype to {torch_dtype} for int4_weight_only quantization, but only bfloat16 is supported right now. Please set the torch_dtype to bfloat16."
+ )
+ if torch_dtype is None:
+ logger.warning_once(
+ "Setting torch_dtype to torch.bfloat16 for int4_weight_only quantization since only bfloat16 is supported right now. Please set torch_dtype=torch.bfloat16 to remove this warning."
+ )
+ torch_dtype = torch.bfloat16
+ if self.quantization_config.quant_type == "int8_dynamic_activation_int8_weight":
+ if torch_dtype is None:
+ logger.info(
+ "Setting torch_dtype to torch.float32 for int8_dynamic_activation_int8_weight quantization as no torch_dtype was specified in from_pretrained"
+ )
+ # we need to set the torch_dtype, otherwise we have dtype mismatch when performing the quantized linear op
+ torch_dtype = torch.float32
+ return torch_dtype
+
+ def adjust_target_dtype(self, torch_dtype: "torch.dtype") -> "torch.dtype":
+ if version.parse(importlib.metadata.version("accelerate")) > version.parse("0.19.0"):
+ from accelerate.utils import CustomDtype
+
+ # Import AOBaseConfig directly since we know we have the right version
+ if self.quantization_config._get_ao_version() > version.Version("0.9.0"):
+ from torchao.core.config import AOBaseConfig
+
+ quant_type = self.quantization_config.quant_type
+ if isinstance(quant_type, AOBaseConfig):
+ # Extract size digit using fuzzy match on the class name
+ config_name = quant_type.__class__.__name__
+ size_digit = fuzzy_match_size(config_name)
+
+ # Map the extracted digit to appropriate dtype
+ if size_digit == "4":
+ return CustomDtype.INT4
+ else:
+ # Default to int8
+ return torch.int8
+
+ # Original mapping for non-AOBaseConfig types
+ map_to_target_dtype = {
+ "int4_weight_only": CustomDtype.INT4,
+ "int8_weight_only": torch.int8,
+ "int8_dynamic_activation_int8_weight": torch.int8,
+ "autoquant": None,
+ }
+ return map_to_target_dtype[self.quantization_config.quant_type]
+ else:
+ raise ValueError(
+ "You are using `device_map='auto'` on a torchao quantized model. To automatically compute"
+ " the appropriate device map, you should upgrade your `accelerate` library with "
+ "`pip install --upgrade accelerate`"
+ )
+
+ def adjust_max_memory(self, max_memory: Dict[str, Union[int, str]]) -> Dict[str, Union[int, str]]:
+ # need more space for the quantization parameters (e.g. scale). Tested with int4 wo and group size = 128
+ max_memory = {key: val * 0.9 for key, val in max_memory.items()}
+ return max_memory
+
+ def _process_model_before_weight_loading(
+ self, model: "PreTrainedModel", keep_in_fp32_modules: Optional[List[str]] = None, **kwargs
+ ):
+ self.modules_to_not_convert = self.get_modules_to_not_convert(
+ model, self.quantization_config.modules_to_not_convert, keep_in_fp32_modules
+ )
+ return
+
+ def check_quantized_param(
+ self,
+ model: "PreTrainedModel",
+ param_value: "torch.Tensor",
+ param_name: str,
+ state_dict: Dict[str, Any],
+ **kwargs,
+ ) -> bool:
+ if self.quantization_config.quant_type == "autoquant":
+ return False
+
+ param_device = kwargs.pop("param_device", None)
+ # check if the param_name is not in self.modules_to_not_convert
+ if any((key + "." in param_name) or (key == param_name) for key in self.modules_to_not_convert):
+ return False
+ elif param_device == "cpu" and self.offload:
+ # We don't quantize weights that we offload
+ return False
+ else:
+ # we only quantize the weight of nn.Linear
+ module, tensor_name = get_module_from_name(model, param_name)
+ return isinstance(module, torch.nn.Linear) and (tensor_name == "weight")
+
+ def create_quantized_param(
+ self,
+ model: "PreTrainedModel",
+ param_value: "torch.Tensor",
+ param_name: str,
+ target_device: "torch.device",
+ state_dict: Dict[str, Any],
+ unexpected_keys: List[str],
+ ):
+ """
+ Each nn.Linear layer that needs to be quantized is processsed here.
+ First, we set the value the weight tensor, then we move it to the target device. Finally, we quantize the module.
+ """
+ if self.quantization_config.quant_type == "autoquant":
+ return
+
+ from torchao.quantization import quantize_
+
+ module, tensor_name = get_module_from_name(model, param_name)
+ if self.pre_quantized:
+ module._parameters[tensor_name] = torch.nn.Parameter(param_value.to(device=target_device))
+ if isinstance(module, nn.Linear):
+ module.extra_repr = types.MethodType(_linear_extra_repr, module)
+ else:
+ assert isinstance(self.quantization_config, TorchAoConfig)
+ module._parameters[tensor_name] = torch.nn.Parameter(param_value).to(device=target_device)
+ quantize_(module, self.quantization_config.get_apply_tensor_subclass())
+
+ def _process_model_after_weight_loading(self, model, **kwargs):
+ """No process required for torchao quantized model"""
+ if self.quantization_config.quant_type == "autoquant":
+ from torchao import autoquant
+ from torchao.quantization import ALL_AUTOQUANT_CLASS_LIST
+
+ model = torch.compile(model, mode="max-autotune")
+ model = autoquant(
+ model,
+ qtensor_class_list=ALL_AUTOQUANT_CLASS_LIST,
+ set_inductor_config=False,
+ **self.quantization_config.quant_type_kwargs,
+ )
+ return model
+ return
+
+ def is_serializable(self, safe_serialization=None) -> bool:
+ if safe_serialization:
+ logger.warning(
+ "torchao quantized model does not support safe serialization, please set `safe_serialization` to False"
+ )
+ return False
+ _is_torchao_serializable = version.parse(importlib.metadata.version("huggingface_hub")) >= version.parse(
+ "0.25.0"
+ )
+ if not _is_torchao_serializable:
+ logger.warning("torchao quantized model is only serializable after huggingface_hub >= 0.25.0 ")
+ if self.offload and self.quantization_config.modules_to_not_convert is None:
+ logger.warning(
+ "The model contains offloaded modules and these modules are not quantized. We don't recommend saving the model as we won't be able to reload them."
+ "If you want to specify modules to not quantize, please specify modules_to_not_convert in the quantization_config."
+ )
+ return False
+ return _is_torchao_serializable
+
+ @property
+ def is_trainable(self) -> bool:
+ supported_quant_types_for_training = [
+ "int8_weight_only",
+ "int8_dynamic_activation_int8_weight",
+ ]
+ return self.quantization_config.quant_type in supported_quant_types_for_training
+
+ @property
+ def is_compileable(self) -> bool:
+ return True
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_vptq.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_vptq.py
new file mode 100644
index 0000000000000000000000000000000000000000..85483357448f15df66dc77b4e0e9158970ab1764
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizer_vptq.py
@@ -0,0 +1,99 @@
+# Copyright 2024 The HuggingFace Inc. 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.
+from typing import TYPE_CHECKING, List, Optional
+
+from .base import HfQuantizer
+
+
+if TYPE_CHECKING:
+ from ..modeling_utils import PreTrainedModel
+
+from ..utils import is_accelerate_available, is_torch_available, is_vptq_available, logging
+from ..utils.quantization_config import QuantizationConfigMixin
+
+
+if is_torch_available():
+ import torch
+
+logger = logging.get_logger(__name__)
+
+
+class VptqHfQuantizer(HfQuantizer):
+ """
+ Quantizer of the VPTQ method. Enables the loading of prequantized models.
+ """
+
+ requires_calibration = True
+ required_packages = ["vptq"]
+
+ def __init__(self, quantization_config: QuantizationConfigMixin, **kwargs):
+ super().__init__(quantization_config, **kwargs)
+ self.quantization_config = quantization_config
+
+ def validate_environment(self, *args, **kwargs):
+ if not is_accelerate_available():
+ raise ImportError("Using `vptq` quantization requires Accelerate: `pip install accelerate`")
+
+ if not is_vptq_available():
+ raise ImportError("Using `vptq` quantization requires VPTQ>=0.0.4: `pip install -U vptq`")
+
+ def update_torch_dtype(self, torch_dtype: "torch.dtype") -> "torch.dtype":
+ if torch_dtype is None:
+ if torch.cuda.is_available():
+ torch_dtype = torch.float16
+ logger.info(
+ "CUDA available. Assuming VPTQ inference on GPU and loading the model in `torch.float16`. To overwrite it, set `torch_dtype` manually."
+ )
+ else:
+ import vptq
+
+ device_availability = getattr(vptq, "device_availability", lambda device: False)
+ if device_availability("cpu") is True:
+ raise RuntimeError("No GPU found. Please wait for the next release of VPTQ to use CPU inference")
+ torch_dtype = torch.float32
+ logger.info("No GPU found. Assuming VPTQ inference on CPU and loading the model in `torch.float32`.")
+ return torch_dtype
+
+ def _process_model_before_weight_loading(
+ self,
+ model: "PreTrainedModel",
+ keep_in_fp32_modules: Optional[List[str]] = None,
+ **kwargs,
+ ):
+ """
+ we don't have param like modules_to_not_convert to indicate which layers should not be quantized
+ because `quantization_config` include the layers that should be quantized
+ """
+ from ..integrations import replace_with_vptq_linear
+
+ self.modules_to_not_convert = self.get_modules_to_not_convert(
+ model, self.quantization_config.modules_to_not_convert, keep_in_fp32_modules
+ )
+
+ replace_with_vptq_linear(
+ model,
+ quantization_config=self.quantization_config,
+ modules_to_not_convert=self.modules_to_not_convert,
+ )
+ model.config.quantization_config = self.quantization_config
+
+ def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs):
+ return model
+
+ @property
+ def is_trainable(self, model: Optional["PreTrainedModel"] = None):
+ return False
+
+ def is_serializable(self, safe_serialization=None):
+ return True
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizers_utils.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizers_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..bbed8317049f9e2530c03dd40730e9c79cb879b2
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quantizers_utils.py
@@ -0,0 +1,21 @@
+# Copyright 2024 The HuggingFace Inc. 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.
+from typing import Any, Tuple
+
+
+def get_module_from_name(module, tensor_name: str) -> Tuple[Any, str]:
+ if "." in tensor_name:
+ module_name, tensor_name = tensor_name.rsplit(".", 1)
+ module = module.get_submodule(module_name)
+ return module, tensor_name
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quanto.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quanto.py
new file mode 100644
index 0000000000000000000000000000000000000000..1c5702321937da5f1238427156452cbccdbcea4a
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/quanto.py
@@ -0,0 +1,98 @@
+# Copyright 2024 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.
+
+from ..utils import is_optimum_quanto_available, is_torch_available, logging
+
+
+if is_torch_available():
+ import torch
+
+logger = logging.get_logger(__name__)
+
+
+def replace_with_quanto_layers(
+ model,
+ quantization_config=None,
+ modules_to_not_convert=None,
+ current_key_name=None,
+ has_been_replaced=False,
+):
+ """
+ Public method that recursively replaces the Linear layers of the given model with Quanto quantized layers.
+ Returns the converted model and a boolean that indicates if the conversion has been successfull or not.
+
+ Args:
+ model (`torch.nn.Module`):
+ The model to convert, can be any `torch.nn.Module` instance.
+ quantization_config (`AqlmConfig`, defaults to `None`):
+ The quantization config object that contains the quantization parameters.
+ modules_to_not_convert (`list`, *optional*, defaults to `None`):
+ A list of modules to not convert. If a module name is in the list (e.g. `lm_head`), it will not be
+ converted.
+ current_key_name (`list`, *optional*, defaults to `None`):
+ A list that contains the current key name. This is used for recursion and should not be passed by the user.
+ has_been_replaced (`bool`, *optional*, defaults to `None`):
+ A boolean that indicates if the conversion has been successful or not. This is used for recursion and
+ should not be passed by the user.
+ """
+ from accelerate import init_empty_weights
+
+ if is_optimum_quanto_available():
+ from optimum.quanto import QLayerNorm, QLinear, qfloat8, qint2, qint4, qint8
+
+ w_mapping = {"float8": qfloat8, "int8": qint8, "int4": qint4, "int2": qint2}
+ a_mapping = {None: None, "float8": qfloat8, "int8": qint8}
+
+ if modules_to_not_convert is None:
+ modules_to_not_convert = []
+
+ for name, module in model.named_children():
+ if current_key_name is None:
+ current_key_name = []
+ current_key_name.append(name)
+
+ if not any(key in ".".join(current_key_name) for key in modules_to_not_convert):
+ with init_empty_weights():
+ if isinstance(module, torch.nn.Linear):
+ model._modules[name] = QLinear(
+ in_features=module.in_features,
+ out_features=module.out_features,
+ bias=module.bias is not None,
+ dtype=module.weight.dtype,
+ weights=w_mapping[quantization_config.weights],
+ activations=a_mapping[quantization_config.activations],
+ )
+ model._modules[name].requires_grad_(False)
+ has_been_replaced = True
+ elif isinstance(module, torch.nn.LayerNorm):
+ if quantization_config.activations is not None:
+ model._modules[name] = QLayerNorm(
+ module.normalized_shape,
+ module.eps,
+ module.elementwise_affine,
+ module.bias is not None,
+ activations=a_mapping[quantization_config.activations],
+ )
+ has_been_replaced = True
+ if len(list(module.children())) > 0:
+ _, has_been_replaced = replace_with_quanto_layers(
+ module,
+ quantization_config=quantization_config,
+ modules_to_not_convert=modules_to_not_convert,
+ current_key_name=current_key_name,
+ has_been_replaced=has_been_replaced,
+ )
+ # Remove the last key for recursion
+ current_key_name.pop(-1)
+ return model, has_been_replaced
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/safetensors_conversion.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/safetensors_conversion.py
new file mode 100644
index 0000000000000000000000000000000000000000..f1612d3ea57c98fd1d383887cfbeb4e2882d3963
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/safetensors_conversion.py
@@ -0,0 +1,105 @@
+from typing import Optional
+
+import requests
+from huggingface_hub import Discussion, HfApi, get_repo_discussions
+
+from .utils import cached_file, http_user_agent, logging
+
+
+logger = logging.get_logger(__name__)
+
+
+def previous_pr(api: HfApi, model_id: str, pr_title: str, token: str) -> Optional["Discussion"]:
+ main_commit = api.list_repo_commits(model_id, token=token)[0].commit_id
+ for discussion in get_repo_discussions(repo_id=model_id, token=token):
+ if discussion.title == pr_title and discussion.status == "open" and discussion.is_pull_request:
+ commits = api.list_repo_commits(model_id, revision=discussion.git_reference, token=token)
+
+ if main_commit == commits[1].commit_id:
+ return discussion
+ return None
+
+
+def spawn_conversion(token: str, private: bool, model_id: str):
+ logger.info("Attempting to convert .bin model on the fly to safetensors.")
+
+ safetensors_convert_space_url = "https://safetensors-convert.hf.space"
+ sse_url = f"{safetensors_convert_space_url}/call/run"
+
+ def start(_sse_connection):
+ for line in _sse_connection.iter_lines():
+ line = line.decode()
+ if line.startswith("event:"):
+ status = line[7:]
+ logger.debug(f"Safetensors conversion status: {status}")
+
+ if status == "complete":
+ return
+ elif status == "heartbeat":
+ logger.debug("Heartbeat")
+ else:
+ logger.debug(f"Unknown status {status}")
+ else:
+ logger.debug(line)
+
+ data = {"data": [model_id, private, token]}
+
+ result = requests.post(sse_url, stream=True, json=data).json()
+ event_id = result["event_id"]
+
+ with requests.get(f"{sse_url}/{event_id}", stream=True) as sse_connection:
+ try:
+ logger.debug("Spawning safetensors automatic conversion.")
+ start(sse_connection)
+ except Exception as e:
+ logger.warning(f"Error during conversion: {repr(e)}")
+
+
+def get_conversion_pr_reference(api: HfApi, model_id: str, **kwargs):
+ private = api.model_info(model_id).private
+
+ logger.info("Attempting to create safetensors variant")
+ pr_title = "Adding `safetensors` variant of this model"
+ token = kwargs.get("token")
+
+ # This looks into the current repo's open PRs to see if a PR for safetensors was already open. If so, it
+ # returns it. It checks that the PR was opened by the bot and not by another user so as to prevent
+ # security breaches.
+ pr = previous_pr(api, model_id, pr_title, token=token)
+
+ if pr is None or (not private and pr.author != "SFconvertbot"):
+ spawn_conversion(token, private, model_id)
+ pr = previous_pr(api, model_id, pr_title, token=token)
+ else:
+ logger.info("Safetensors PR exists")
+
+ sha = f"refs/pr/{pr.num}"
+
+ return sha
+
+
+def auto_conversion(pretrained_model_name_or_path: str, ignore_errors_during_conversion=False, **cached_file_kwargs):
+ try:
+ api = HfApi(token=cached_file_kwargs.get("token"), headers={"user-agent": http_user_agent()})
+ sha = get_conversion_pr_reference(api, pretrained_model_name_or_path, **cached_file_kwargs)
+
+ if sha is None:
+ return None, None
+ cached_file_kwargs["revision"] = sha
+ del cached_file_kwargs["_commit_hash"]
+
+ # This is an additional HEAD call that could be removed if we could infer sharded/non-sharded from the PR
+ # description.
+ sharded = api.file_exists(
+ pretrained_model_name_or_path,
+ "model.safetensors.index.json",
+ revision=sha,
+ token=cached_file_kwargs.get("token"),
+ )
+ filename = "model.safetensors.index.json" if sharded else "model.safetensors"
+
+ resolved_archive_file = cached_file(pretrained_model_name_or_path, filename, **cached_file_kwargs)
+ return resolved_archive_file, sha, sharded
+ except Exception as e:
+ if not ignore_errors_during_conversion:
+ raise e
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/sdpa_attention.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/sdpa_attention.py
new file mode 100644
index 0000000000000000000000000000000000000000..9c924c048ad52929a2d0f890a22295d5a56ef505
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/sdpa_attention.py
@@ -0,0 +1,65 @@
+from typing import Optional, Tuple
+
+import torch
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+ """
+ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+ num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+ """
+ batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+ if n_rep == 1:
+ return hidden_states
+ hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+ return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+def sdpa_attention_forward(
+ module: torch.nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: Optional[torch.Tensor],
+ dropout: float = 0.0,
+ scaling: Optional[float] = None,
+ is_causal: Optional[bool] = None,
+ **kwargs,
+) -> Tuple[torch.Tensor, None]:
+ if hasattr(module, "num_key_value_groups"):
+ key = repeat_kv(key, module.num_key_value_groups)
+ value = repeat_kv(value, module.num_key_value_groups)
+
+ causal_mask = attention_mask
+ if attention_mask is not None and causal_mask.ndim == 4:
+ causal_mask = causal_mask[:, :, :, : key.shape[-2]]
+
+ # SDPA with memory-efficient backend is bugged with non-contiguous inputs and custom attn_mask for some torch versions
+ # Reference: https://github.com/pytorch/pytorch/issues/112577.
+ query = query.contiguous()
+ key = key.contiguous()
+ value = value.contiguous()
+
+ # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment
+ # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling.
+ # Note that it is important to check first for the shape, otherwise compile will fail with `argument 'is_causal' must be bool, not SymBool`
+ if is_causal is None:
+ is_causal = query.shape[2] > 1 and causal_mask is None
+
+ # Shapes (e.g. query.shape[2]) are tensors during jit tracing, resulting in `is_causal` being a tensor.
+ # We convert it to a bool for the SDPA kernel that only accepts bools.
+ if torch.jit.is_tracing() and isinstance(is_causal, torch.Tensor):
+ is_causal = is_causal.item()
+
+ attn_output = torch.nn.functional.scaled_dot_product_attention(
+ query,
+ key,
+ value,
+ attn_mask=causal_mask,
+ dropout_p=dropout,
+ scale=scaling,
+ is_causal=is_causal,
+ )
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, None
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/special_tokens_map.json b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/special_tokens_map.json
new file mode 100644
index 0000000000000000000000000000000000000000..ac23c0aaa2434523c494330aeb79c58395378103
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/special_tokens_map.json
@@ -0,0 +1,31 @@
+{
+ "additional_special_tokens": [
+ "<|im_start|>",
+ "<|im_end|>",
+ "<|object_ref_start|>",
+ "<|object_ref_end|>",
+ "<|box_start|>",
+ "<|box_end|>",
+ "<|quad_start|>",
+ "<|quad_end|>",
+ "<|vision_start|>",
+ "<|vision_end|>",
+ "<|vision_pad|>",
+ "<|image_pad|>",
+ "<|video_pad|>"
+ ],
+ "eos_token": {
+ "content": "<|im_end|>",
+ "lstrip": false,
+ "normalized": false,
+ "rstrip": false,
+ "single_word": false
+ },
+ "pad_token": {
+ "content": "<|endoftext|>",
+ "lstrip": false,
+ "normalized": false,
+ "rstrip": false,
+ "single_word": false
+ }
+}
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/spqr.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/spqr.py
new file mode 100644
index 0000000000000000000000000000000000000000..58b71740d37c77f24fe60ff65fe179b2737c4aa7
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/spqr.py
@@ -0,0 +1,122 @@
+# Copyright 2024 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.
+"SpQR (Sparse-Quantized Representation) integration file"
+
+from ..utils import is_accelerate_available, is_spqr_available, is_torch_available
+
+
+if is_torch_available():
+ import torch.nn as nn
+
+
+def replace_with_spqr_linear(
+ model,
+ quantization_config=None,
+ modules_to_not_convert=None,
+ current_key_name=None,
+ has_been_replaced=False,
+):
+ """
+ Public method that recursively replaces the Linear layers of the given model with SpQR quantized layers.
+ `accelerate` is needed to use this method. Returns the converted model and a boolean that indicates if the
+ conversion has been successful or not.
+
+ Args:
+ model (`torch.nn.Module`):
+ The model to convert, can be any `torch.nn.Module` instance.
+ quantization_config (`SpQRConfig`):
+ The quantization config object that contains the quantization parameters.
+ modules_to_not_convert (`list[str]`, *optional*):
+ A list of nn.Linear weights to not convert. If a parameter path is in the list (e.g. `lm_head.weight`), the corresponding module will not be
+ converted.
+ current_key_name (`list`, *optional*):
+ A list that contains the current key name. This is used for recursion and should not be passed by the user.
+ has_been_replaced (`bool`, *optional*):
+ A boolean that indicates if the conversion has been successful or not. This is used for recursion and
+ should not be passed by the user.
+ """
+ if modules_to_not_convert is None:
+ modules_to_not_convert = []
+
+ if is_accelerate_available():
+ from accelerate import init_empty_weights
+ if is_spqr_available():
+ from spqr_quant import QuantizedLinear
+
+ for name, module in model.named_children():
+ if current_key_name is None:
+ current_key_name = []
+ current_key_name.append(name)
+
+ if isinstance(module, nn.Linear):
+ # Check if the current key is not in the `modules_to_not_convert`
+ if ".".join(current_key_name) + ".weight" not in modules_to_not_convert:
+ with init_empty_weights():
+ tensor_name = ".".join(current_key_name)
+
+ shapes = quantization_config.shapes
+ shapes_keys = shapes.keys()
+
+ shapes_valid = (
+ f"{tensor_name}.dense_weights.shape" in shapes_keys
+ and f"{tensor_name}.row_offsets.shape" in shapes_keys
+ and f"{tensor_name}.col_vals.shape" in shapes_keys
+ and f"{tensor_name}.in_perm.shape" in shapes_keys
+ )
+
+ if not shapes_valid:
+ raise ValueError(
+ f"The SpQR quantization config does not contain the shape "
+ f"configuration for {tensor_name}. This indicates that the "
+ f"configuration is either invalid or corrupted."
+ )
+
+ dense_weights_shape = shapes[f"{tensor_name}.dense_weights.shape"]
+ row_offsets_shape = shapes[f"{tensor_name}.row_offsets.shape"]
+ col_vals_shape = shapes[f"{tensor_name}.col_vals.shape"]
+ in_perm_shape = shapes[f"{tensor_name}.in_perm.shape"]
+
+ in_features = module.in_features
+ out_features = module.out_features
+
+ model._modules[name] = QuantizedLinear.create_placehodler(
+ rows=out_features,
+ cols=in_features,
+ bits=quantization_config.bits,
+ beta1=quantization_config.beta1,
+ beta2=quantization_config.beta2,
+ dense_weights_shape=dense_weights_shape,
+ row_offsets_shape=row_offsets_shape,
+ col_vals_shape=col_vals_shape,
+ in_perm_shape=in_perm_shape,
+ )
+ has_been_replaced = True
+
+ # Store the module class in case we need to transpose the weight later
+ model._modules[name].source_cls = type(module)
+ # Force requires grad to False to avoid unexpected errors
+ model._modules[name].requires_grad_(False)
+ else:
+ pass
+ if len(list(module.children())) > 0:
+ _, has_been_replaced = replace_with_spqr_linear(
+ module,
+ quantization_config=quantization_config,
+ modules_to_not_convert=modules_to_not_convert,
+ current_key_name=current_key_name,
+ has_been_replaced=has_been_replaced,
+ )
+ # Remove the last key for recursion
+ current_key_name.pop(-1)
+ return model, has_been_replaced
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/stopping_criteria.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/stopping_criteria.py
new file mode 100644
index 0000000000000000000000000000000000000000..4627aeb9702707602d1f09e2f2121ee47338fe39
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/stopping_criteria.py
@@ -0,0 +1,519 @@
+import time
+import warnings
+from abc import ABC
+from collections import OrderedDict
+from copy import deepcopy
+from typing import Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+from torch.nn import functional as F
+
+from ..pytorch_utils import isin_mps_friendly
+from ..tokenization_utils_base import PreTrainedTokenizerBase
+from ..utils import add_start_docstrings, logging
+
+
+logger = logging.get_logger(__name__)
+# We maintain a module-level cache of the embedding vectors for the stop string criterion
+# because they are slow to compute
+STOP_STRING_EMBEDDING_CACHE = OrderedDict()
+
+
+STOPPING_CRITERIA_INPUTS_DOCSTRING = r"""
+ Args:
+ input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+ Indices of input sequence tokens in the vocabulary.
+
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+ [`PreTrainedTokenizer.__call__`] for details.
+
+ [What are input IDs?](../glossary#input-ids)
+ scores (`torch.FloatTensor` of shape `(batch_size, config.vocab_size)`):
+ Prediction scores of a language modeling head. These can be scores for each vocabulary token before SoftMax
+ or scores for each vocabulary token after SoftMax. If this stopping criteria depends on the `scores` input,
+ make sure you pass `return_dict_in_generate=True, output_scores=True` to `generate`.
+ kwargs (`Dict[str, Any]`, *optional*):
+ Additional stopping criteria specific kwargs.
+
+ Return:
+ `torch.BoolTensor`. (`torch.BoolTensor` of shape `(batch_size, 1)`), where `True` indicates we stop generation
+ for a particular row, `True` indicates we should continue.
+
+"""
+
+
+class StoppingCriteria(ABC):
+ """Abstract base class for all stopping criteria that can be applied during generation.
+
+ If your stopping criteria depends on the `scores` input, make sure you pass `return_dict_in_generate=True,
+ output_scores=True` to `generate`.
+ """
+
+ @add_start_docstrings(STOPPING_CRITERIA_INPUTS_DOCSTRING)
+ def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> torch.BoolTensor:
+ raise NotImplementedError("StoppingCriteria needs to be subclassed")
+
+
+class MaxLengthCriteria(StoppingCriteria):
+ """
+ This class can be used to stop generation whenever the full generated number of tokens exceeds `max_length`. Keep
+ in mind for decoder-only type of transformers, this will include the initial prompted tokens.
+
+ Args:
+ max_length (`int`):
+ The maximum length that the output sequence can have in number of tokens.
+ max_position_embeddings (`int`, *optional*):
+ The maximum model length, as defined by the model's `config.max_position_embeddings` attribute.
+ """
+
+ def __init__(self, max_length: int, max_position_embeddings: Optional[int] = None):
+ self.max_length = max_length
+ self.max_position_embeddings = max_position_embeddings
+
+ @add_start_docstrings(STOPPING_CRITERIA_INPUTS_DOCSTRING)
+ def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> torch.BoolTensor:
+ cur_len = input_ids.shape[-1]
+ is_done = cur_len >= self.max_length
+ if self.max_position_embeddings is not None and not is_done and cur_len >= self.max_position_embeddings:
+ logger.warning_once(
+ "This is a friendly reminder - the current text generation call will exceed the model's predefined "
+ f"maximum length ({self.max_position_embeddings}). Depending on the model, you may observe "
+ "exceptions, performance degradation, or nothing at all."
+ )
+ return torch.full((input_ids.shape[0],), is_done, device=input_ids.device, dtype=torch.bool)
+
+
+class MaxTimeCriteria(StoppingCriteria):
+ """
+ This class can be used to stop generation whenever the full generation exceeds some amount of time. By default, the
+ time will start being counted when you initialize this function. You can override this by passing an
+ `initial_time`.
+
+ Args:
+ max_time (`float`):
+ The maximum allowed time in seconds for the generation.
+ initial_time (`float`, *optional*, defaults to `time.time()`):
+ The start of the generation allowed time.
+ """
+
+ def __init__(self, max_time: float, initial_timestamp: Optional[float] = None):
+ self.max_time = max_time
+ self.initial_timestamp = time.time() if initial_timestamp is None else initial_timestamp
+
+ @add_start_docstrings(STOPPING_CRITERIA_INPUTS_DOCSTRING)
+ def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> torch.BoolTensor:
+ is_done = time.time() - self.initial_timestamp > self.max_time
+ return torch.full((input_ids.shape[0],), is_done, device=input_ids.device, dtype=torch.bool)
+
+
+class StopStringCriteria(StoppingCriteria):
+ """
+ This class can be used to stop generation whenever specific string sequences are generated. It preprocesses
+ the strings together with the tokenizer vocab to find positions where tokens can validly complete the stop strings.
+
+ Generation is stopped as soon as a token is generated that completes any of the stop strings.
+ We want to catch any instance in which the stop string would be present in the decoded output, which means
+ we must also catch cases with "overhangs" off one or both ends. To make this more concrete, for the stop string
+ "stop", any of the following token sequences would trigger the match:
+
+ - ["st", "op"]
+ - ["stop"]
+ - ["st", "opera"]
+ - ["sto", "pper"]
+ - ["las", "topper"]
+ - ["s", "to", "pped"]
+
+ Note that a match will only be triggered if the stop string is at the end of the generated sequence. In other
+ words, these sequences will not trigger a match:
+
+ - ["stop", "at"]
+ - ["st", "op", "at"]
+ - ["st", "opera", "tion"]
+
+ The reason these are not a match is that the stop string does not overlap with the final token. If you can remove
+ one or more tokens from the end of the sequence without destroying the stop string, then this criterion will not
+ match that stop string. This is by design; because this check is run after each token is generated, we can't miss a
+ valid stop string if one is generated, but we don't want to halt generation just because the stop string exists
+ somewhere in the past input_ids.
+
+ How is the match actually performed, though? We do it in quite a confusing way, because we want the entire match
+ process to be compilable with Torch or XLA, which means we cannot use standard string methods. However, it is possible,
+ with some work, to do string matching with pure tensor operations. We'll begin by describing the algorithm we use
+ with standard string operations, and then at the end we'll explain how this is converted to pure tensor operations.
+
+ The key to the algorithm is an observation: Because the stop string must overlap with the end of the token sequence, we can start at
+ the end of the sequence and work backwards. Specifically, we check that there is an overlap between the start of
+ the final token and the end of the stop_string, or to put it another way, stop_string[-i:] == token[:i] for
+ some i > 0. If you look at the positive examples above, you'll see the last token in all of them fulfills this
+ property:
+
+ - ["st", "op"] (overlap is "op", overlap length == 2)
+ - ["stop"] (overlap is "stop", overlap length == 4)
+ - ["st", "opera"] (overlap is "op", overlap length == 2)
+ - ["sto", "pper"] (overlap is "p", overlap length == 1)
+ - ["las", "topper"] (overlap is "top", overlap length == 3)
+ - ["s", "to", "pped"] (overlap is "p", overlap length == 1)
+
+ It's impossible to construct a matching sequence that does not have this property (feel free to verify this
+ yourself). However, although this overlap between the start of the final token and the end of the stop string is
+ necessary for a match, it is not sufficient. We also need to check that the rest of the token sequence is
+ consistent with the stop string.
+
+ How do we do that? Let's use ["s", "to", "pped"] as an example. We know that the final token, "pped", has an
+ overlap of 1 with the stop string, "stop". We then go back to the previous token, "to". Since we have already
+ matched 1 character from the stop string, the remainder to check is "sto". We check that the next token "to"
+ matches the end of the remainder, which it does. We have now matched 3 characters from the stop string, and the
+ remainder to match is "s". We go back to the previous token again, which is also "s". This is a match, and so
+ we have matched the entire stop string.
+
+ How does it work when the tokens run off the start of the stop string, though? Let's consider the example of
+ ["las", "topper"]. The final token, "topper", has an overlap of 3 with the stop string, "stop". Therefore,
+ the remaining stop string to match is "s". We go back to the previous token, "las". Because the remainder to
+ match is just "s", with length 1, we consider only the final 1 character from the token, which is "s". This
+ matches the stop string, and so the entire string is matched.
+
+ How do we compute these matches with tensor operations, though? Simply: we efficiently precompute the necessary
+ information for all tokens! For every token, we compute:
+ - Its overlap with the end of the stop string, if any
+ - The positions inside the stop string where the token matches, including matches that run off the start.
+ - The total length of the token
+
+ For example, for the token "pped", we would compute an end overlap of 1, no internal matching positions,
+ and a length of 4. For the token "to", we would compute no end overlap, a single internal matching position
+ of 1 (counting from the end), and a length of 2. For the token "s", we would compute no end overlap,
+ a single internal matching position of 3 (again counting from the end) and a length of 1.
+
+ As long as we have this information, we can execute the algorithm above without any string comparison
+ operations. We simply perform the following steps:
+ - Check if the final token has an end-overlap with the start string
+ - Continue backwards, keeping track of how much of the stop string we've matched so far
+ - At each point, check if the next token has the current position as one of its valid positions
+ - Continue until either a match fails, or we completely match the whole stop string
+
+ Again, consider ["s", "to", "pped"] as an example. "pped" has an end overlap of 1, so we can begin a match.
+ We have matched 1 character so far, so we check that the next token "to", has 1 as a valid position (again,
+ counting from the end). It does, so we add the length of "to" to our position tracker. We have now matched
+ 3 characters, so we check that the next token "s" has 3 as a valid position. It does, so we add its length
+ to the position tracker. The position tracker is now 4, which is the length of the stop string. We have matched the
+ entire stop string.
+
+ In the second case, ["las", "topper"], "topper" has an end overlap of 3, so we can begin a match. We have
+ matched 3 characters so far, so we check that the next token "las" has 3 as a valid position. It does, because we
+ allow tokens to match positions that run off the start of the stop string. We add its length to the position
+ tracker. The position tracker is now 6, which is greater than the length of the stop string! Don't panic, though -
+ this also counts as a match of the stop string. We have matched the entire stop string.
+
+
+ Args:
+ tokenizer (`PreTrainedTokenizer`):
+ The model's associated tokenizer (necessary to extract vocab and tokenize the termination sequences)
+ stop_strings (`Union[str, List[str]]`):
+ A list of strings that should end generation. If a string is passed, it will be treated like a
+ list with a single element.
+
+ Examples:
+
+ ```python
+ >>> from transformers import AutoModelForCausalLM, AutoTokenizer
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2")
+ >>> model = AutoModelForCausalLM.from_pretrained("microsoft/phi-2")
+ >>> inputs = tokenizer("The biggest states in the USA by land area:", return_tensors="pt")
+
+ >>> gen_out = model.generate(**inputs)
+ >>> print(tokenizer.batch_decode(gen_out, skip_special_tokens=True)[0])
+ The biggest states in the USA by land area:
+ - Alaska
+ - Texas
+ - California
+
+ >>> # Passing one or more stop strings will halt generation after those strings are emitted
+ >>> # Note that generating with stop strings requires you to pass the tokenizer too
+ >>> gen_out = model.generate(**inputs, stop_strings=["Texas"], tokenizer=tokenizer)
+ >>> print(tokenizer.batch_decode(gen_out, skip_special_tokens=True)[0])
+ The biggest states in the USA by land area:
+ - Alaska
+ - Texas
+ ```
+ """
+
+ def __init__(self, tokenizer: PreTrainedTokenizerBase, stop_strings: Union[str, List[str]]):
+ if isinstance(stop_strings, str):
+ stop_strings = [stop_strings]
+ self.stop_strings: Tuple[str, ...] = tuple(stop_strings)
+ vocab = tokenizer.get_vocab()
+ token_list, token_indices = tuple(vocab.keys()), tuple(vocab.values())
+ self.embedding_vec, self.max_valid_positions, self.max_valid_end_lens = self.clean_and_embed_tokens_with_cache(
+ token_list, token_indices, tokenizer
+ )
+
+ self.maximum_token_len = max([len(stop_string) for stop_string in self.stop_strings])
+ self.num_stop_strings = len(self.stop_strings)
+ self.target_lens = torch.tensor([len(stop_string) for stop_string in stop_strings], dtype=torch.int32)
+
+ def clean_and_embed_tokens_with_cache(self, token_list, token_indices, tokenizer):
+ # We don't use the tokenizer in the cache key, because I don't trust it to have well-behaved equality
+ if (token_list, token_indices, self.stop_strings) in STOP_STRING_EMBEDDING_CACHE:
+ embedding_vec, max_valid_positions, max_valid_end_lens = STOP_STRING_EMBEDDING_CACHE[
+ (token_list, token_indices, self.stop_strings)
+ ]
+ STOP_STRING_EMBEDDING_CACHE.move_to_end((token_list, token_indices, self.stop_strings))
+ else:
+ clean_token_list, clean_token_indices = self.clean_tokenizer_vocab(tokenizer)
+ embedding_vec, max_valid_positions, max_valid_end_lens = self._stop_string_create_embedding_vec(
+ clean_token_list, clean_token_indices, self.stop_strings
+ )
+ STOP_STRING_EMBEDDING_CACHE[(token_list, token_indices, self.stop_strings)] = (
+ embedding_vec,
+ max_valid_positions,
+ max_valid_end_lens,
+ )
+ if len(STOP_STRING_EMBEDDING_CACHE) > 8:
+ STOP_STRING_EMBEDDING_CACHE.popitem(last=False) # Pop from the start, the least recently used item
+ return embedding_vec, max_valid_positions, max_valid_end_lens
+
+ @staticmethod
+ def clean_tokenizer_vocab(tokenizer, static_prefix="abcdef"):
+ """
+ This method turns a tokenizer vocab into a "clean" vocab where each token represents the actual string
+ it will yield, without any special prefixes like "##" or "Ġ". This is trickier than it looks - the method
+ tokenizer.convert_tokens_to_string() does not always return the correct string because of issues with prefix
+ space addition/removal. To work around this, we add a static prefix to the start of the token, then remove
+ it (and any prefix that may have been introduced with it) after calling convert_tokens_to_string().
+ """
+ vocab = tokenizer.get_vocab()
+ clean_token_list = []
+ clean_token_indices = []
+ sentence_base = tokenizer(static_prefix, add_special_tokens=False)["input_ids"]
+ tokens_base = [tokenizer._convert_id_to_token(tok) for tok in sentence_base]
+ for token, token_idx in vocab.items():
+ token_string = tokenizer.convert_tokens_to_string(tokens_base + [token])
+ token_string = token_string[token_string.index(static_prefix) + len(static_prefix) :]
+ clean_token_list.append(token_string)
+ clean_token_indices.append(token_idx)
+ return tuple(clean_token_list), tuple(clean_token_indices)
+
+ @staticmethod
+ def _stop_string_get_matching_positions(
+ token_list, token_indices, stop_strings
+ ) -> Tuple[Dict[str, Dict[str, List[int]]], Dict[str, Dict[str, List[int]]]]:
+ """This function preprocesses stop strings and the tokenizer vocabulary to determine where tokens can
+ validly appear in the stop strings. For each token, it computes a list of positions in the stop string where the
+ token appears, as well as a list of the possible "end overlaps" for that token - that is, the number of characters
+ from the end of the stop string that overlap with the start of the token, which can have more than one value.
+
+ The reason for computing these may seem a bit cryptic - please see the docstring for StopStringCriteria for a full
+ explanation of what these values are for!"""
+
+ token_valid_positions = {}
+ token_end_overlaps = {}
+ for stop_string in stop_strings:
+ reversed_stop_string = stop_string[::-1]
+ token_valid_positions[stop_string] = {}
+ token_end_overlaps[stop_string] = {}
+ for token, tok_idx in zip(token_list, token_indices):
+ reversed_token = token[::-1]
+ matching_positions = []
+ possible_end_lengths = []
+ for i in range(1 - len(token), len(stop_string)):
+ if i < 0:
+ tok = reversed_token[-i:]
+ i = 0
+ else:
+ tok = reversed_token
+ stop = reversed_stop_string[i : i + len(tok)]
+ if tok.startswith(stop):
+ if i == 0:
+ possible_end_lengths.append(min(len(tok), len(stop)))
+ else:
+ matching_positions.append(i)
+
+ if matching_positions:
+ token_valid_positions[stop_string][tok_idx] = matching_positions
+ if possible_end_lengths:
+ token_end_overlaps[stop_string][tok_idx] = possible_end_lengths
+ return token_valid_positions, token_end_overlaps
+
+ @staticmethod
+ def _stop_string_create_embedding_vec(token_list, token_indices, stop_strings) -> Dict[str, torch.tensor]:
+ """This function precomputes everything needed for the run-time checks in StopStringCriteria, and packs
+ them into an embedding tensor that can be accessed with pure tensor operations. For the specifics of the values
+ that are precomputed and what they are used for, please refer to the StopStringCriteria docstring!"""
+ token_valid_positions, token_end_overlaps = StopStringCriteria._stop_string_get_matching_positions(
+ token_list, token_indices, stop_strings
+ )
+ all_valid_positions = [len(val) for positions in token_valid_positions.values() for val in positions.values()]
+ # In some cases, tokens may have no valid internal positions (such as single-character stop strings), so
+ # we need a fallback to handle this case
+ max_valid_positions = max(all_valid_positions) if all_valid_positions else 1
+ # There should always be at least one valid end_len, however, so no fallback needed here
+ valid_end_lens = [len(val) for positions in token_end_overlaps.values() for val in positions.values()]
+ if not valid_end_lens:
+ raise ValueError(
+ "Stop string preprocessing was unable to identify tokens matching one or more of the "
+ "supplied stop string(s). This is most often caused by the stop "
+ "strings containing unusual characters that are not in the tokenizer vocabulary."
+ )
+ max_valid_end_lens = max(valid_end_lens)
+ vec_size = len(stop_strings) * (max_valid_positions + max_valid_end_lens) + 1
+ # We use +2 instead of +1 so we can have a dummy entry at the end. We will clamp all token values
+ # over the max to this, ensuring they do not contribute to stop string matching.
+ gather_vec = np.full((max(token_indices) + 2, vec_size), dtype=np.int32, fill_value=-1)
+
+ for i, stop_string in enumerate(stop_strings):
+ positions = token_valid_positions[stop_string]
+ end_lens = token_end_overlaps[stop_string]
+
+ # Since this is lots of very small assignments of lists, we build it with numpy rather
+ # than torch for speed + simplicity, then convert to torch at the end
+ for token_idx, valid_positions in positions.items():
+ gather_vec[token_idx, max_valid_positions * i : max_valid_positions * i + len(valid_positions)] = (
+ valid_positions
+ )
+ for token_idx, possible_end_lens in end_lens.items():
+ gather_vec[
+ token_idx,
+ max_valid_positions * len(stop_strings) + max_valid_end_lens * i : max_valid_positions
+ * len(stop_strings)
+ + max_valid_end_lens * i
+ + len(possible_end_lens),
+ ] = possible_end_lens
+ for token, token_idx in zip(token_list, token_indices):
+ gather_vec[token_idx, -1] = len(token)
+
+ gather_vec = torch.tensor(gather_vec, dtype=torch.int32)
+
+ return gather_vec, max_valid_positions, max_valid_end_lens
+
+ @add_start_docstrings(STOPPING_CRITERIA_INPUTS_DOCSTRING)
+ def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> torch.Tensor:
+ self.embedding_vec = self.embedding_vec.to(input_ids.device)
+ self.target_lens = self.target_lens.to(input_ids.device)
+ # The maximum length we need to consider is 1 token per character. Note that input_ids can also be
+ # *shorter* than the global max, and the code below should be ready for that
+ input_ids = input_ids[:, -self.maximum_token_len :]
+
+ # Flip input_ids because we're only matching strings at the end of the generated sequence
+ flipped_ids = torch.flip(input_ids, (1,))
+
+ # Clip out-of-vocab values to the dummy value at the end of the embedding vector
+ flipped_ids = torch.clamp(flipped_ids, max=self.embedding_vec.size(0) - 1)
+
+ # Size of the vector of positions a single token can match
+ max_valid_positions = self.max_valid_positions
+
+ # The embedding vec contains the valid positions, end_lengths and total lengths for each token
+ embedded = F.embedding(flipped_ids, self.embedding_vec)
+
+ # Now we split the embedding vector. valid_positions is the positions in the stop string the token can fit
+ valid_positions = embedded[:, 1:, : max_valid_positions * self.num_stop_strings].unflatten(
+ -1, (self.num_stop_strings, -1)
+ )
+ # end_lengths is the number of characters from the string, counting from the end, that the token
+ # contains. It can have multiple values if the same token can overlap different end lengths
+ end_lengths = embedded[:, :1, max_valid_positions * self.num_stop_strings : -1].unflatten(
+ -1, (self.num_stop_strings, -1)
+ )
+ # Lengths is the total length of each token. Unlike the others, it always has a single value
+ lengths = embedded[:, 1:, None, -1:] # Insert a dummy dimension for stop_strings even though lengths are const
+
+ # Concatenate lengths onto each possible end_lengths value
+ lengths = lengths.expand((-1, -1, end_lengths.shape[-2], end_lengths.shape[-1]))
+ lengths_with_ends = torch.cat([end_lengths, lengths], dim=1)
+
+ # cumsum() to get the number of matched characters in the stop string after each token
+ cumsum = lengths_with_ends.cumsum(dim=1) # B x maximum_token_len x num_stop_strings x max_valid_end_lens
+
+ # The calculation above assumes that all tokens are in valid positions. Now we mask the ones that are not.
+ # First, tokens match the start of the string if they have a positive value in the end_lengths vector
+ initial_match = end_lengths > 0
+
+ # Tokens continue the string if the cumsum() so far is one of the valid positions for that token
+ # Note that we're actually tracking one cumsum() for for each possible end_length
+ later_match = torch.any(cumsum[:, :-1, :, None] == valid_positions[:, :, :, :, None], axis=-2)
+
+ # The match vector is a boolean vector that indicates which positions have valid tokens
+ match = torch.cat([initial_match, later_match], dim=1)
+
+ # Once a single position does not match, all positions following that position are masked
+ mask = (~match).cumsum(dim=1, dtype=torch.int32)
+ mask = mask == 0
+
+ # The string is matched if we reached a cumsum equal to or greater than the length of the string
+ # before hitting the mask
+ string_matches = torch.amax(cumsum * mask, dim=(1, -1)) >= self.target_lens[None, :]
+
+ # We return a per-sample vector that is True if any stop string is matched for that sample
+ return torch.any(string_matches, dim=-1)
+
+
+class EosTokenCriteria(StoppingCriteria):
+ """
+ This class can be used to stop generation whenever the "end-of-sequence" token is generated.
+ By default, it uses the `model.generation_config.eos_token_id`.
+
+ Args:
+ eos_token_id (`Union[int, List[int], torch.Tensor]`):
+ The id(s) of the *end-of-sequence* token.
+ """
+
+ def __init__(self, eos_token_id: Union[int, List[int], torch.Tensor]):
+ if not isinstance(eos_token_id, torch.Tensor):
+ if isinstance(eos_token_id, int):
+ eos_token_id = [eos_token_id]
+ eos_token_id = torch.tensor(eos_token_id)
+ self.eos_token_id = eos_token_id
+
+ @add_start_docstrings(STOPPING_CRITERIA_INPUTS_DOCSTRING)
+ def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> torch.BoolTensor:
+ self.eos_token_id = self.eos_token_id.to(input_ids.device)
+ is_done = isin_mps_friendly(input_ids[:, -1], self.eos_token_id)
+ return is_done
+
+
+class ConfidenceCriteria(StoppingCriteria):
+ """
+ This class can be used to stop generation whenever assistant model's confidence in its prediction for the current token is lower than the threshold
+ `model.generation_config.assistant_confidence_threshold` even if the number of speculative tokens (defined by `num_assistant_tokens`) is not yet reached.
+
+ Args:
+ assistant_confidence_threshold (`float`):
+ The value of the threshold.
+ """
+
+ def __init__(self, assistant_confidence_threshold):
+ self.assistant_confidence_threshold = assistant_confidence_threshold
+
+ def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> torch.BoolTensor:
+ probs = scores[-1].softmax(-1)
+ p = probs[0, input_ids[0, -1]].item()
+ if p < self.assistant_confidence_threshold:
+ return True
+ return False
+
+
+class StoppingCriteriaList(list):
+ @add_start_docstrings(STOPPING_CRITERIA_INPUTS_DOCSTRING)
+ def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> torch.BoolTensor:
+ is_done = torch.full((input_ids.shape[0],), False, device=input_ids.device, dtype=torch.bool)
+ for criteria in self:
+ is_done = is_done | criteria(input_ids, scores, **kwargs)
+ return is_done
+
+ @property
+ def max_length(self) -> Optional[int]:
+ for stopping_criterium in self:
+ if isinstance(stopping_criterium, MaxLengthCriteria):
+ return stopping_criterium.max_length
+ return None
+
+
+def validate_stopping_criteria(stopping_criteria: StoppingCriteriaList, max_length: int) -> StoppingCriteriaList:
+ stopping_max_length = stopping_criteria.max_length
+ new_stopping_criteria = deepcopy(stopping_criteria)
+ if stopping_max_length is not None and stopping_max_length != max_length:
+ warnings.warn("You set different `max_length` for stopping criteria and `max_length` parameter", UserWarning)
+ elif stopping_max_length is None:
+ new_stopping_criteria.append(MaxLengthCriteria(max_length=max_length))
+ return new_stopping_criteria
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/streamers.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/streamers.py
new file mode 100644
index 0000000000000000000000000000000000000000..c0fc26442b466fdaa09aa7c52e607a37dd0d78ab
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/streamers.py
@@ -0,0 +1,318 @@
+# coding=utf-8
+# Copyright 2023 The HuggingFace Inc. 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.
+
+from __future__ import annotations
+
+import asyncio
+from queue import Queue
+from typing import TYPE_CHECKING, Optional
+
+
+if TYPE_CHECKING:
+ from ..models.auto import AutoTokenizer
+
+
+class BaseStreamer:
+ """
+ Base class from which `.generate()` streamers should inherit.
+ """
+
+ def put(self, value):
+ """Function that is called by `.generate()` to push new tokens"""
+ raise NotImplementedError()
+
+ def end(self):
+ """Function that is called by `.generate()` to signal the end of generation"""
+ raise NotImplementedError()
+
+
+class TextStreamer(BaseStreamer):
+ """
+ Simple text streamer that prints the token(s) to stdout as soon as entire words are formed.
+
+
+
+ The API for the streamer classes is still under development and may change in the future.
+
+
+
+ Parameters:
+ tokenizer (`AutoTokenizer`):
+ The tokenized used to decode the tokens.
+ skip_prompt (`bool`, *optional*, defaults to `False`):
+ Whether to skip the prompt to `.generate()` or not. Useful e.g. for chatbots.
+ decode_kwargs (`dict`, *optional*):
+ Additional keyword arguments to pass to the tokenizer's `decode` method.
+
+ Examples:
+
+ ```python
+ >>> from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
+
+ >>> tok = AutoTokenizer.from_pretrained("openai-community/gpt2")
+ >>> model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2")
+ >>> inputs = tok(["An increasing sequence: one,"], return_tensors="pt")
+ >>> streamer = TextStreamer(tok)
+
+ >>> # Despite returning the usual output, the streamer will also print the generated text to stdout.
+ >>> _ = model.generate(**inputs, streamer=streamer, max_new_tokens=20)
+ An increasing sequence: one, two, three, four, five, six, seven, eight, nine, ten, eleven,
+ ```
+ """
+
+ def __init__(self, tokenizer: "AutoTokenizer", skip_prompt: bool = False, **decode_kwargs):
+ self.tokenizer = tokenizer
+ self.skip_prompt = skip_prompt
+ self.decode_kwargs = decode_kwargs
+
+ # variables used in the streaming process
+ self.token_cache = []
+ self.print_len = 0
+ self.next_tokens_are_prompt = True
+
+ def put(self, value):
+ """
+ Receives tokens, decodes them, and prints them to stdout as soon as they form entire words.
+ """
+ if len(value.shape) > 1 and value.shape[0] > 1:
+ raise ValueError("TextStreamer only supports batch size 1")
+ elif len(value.shape) > 1:
+ value = value[0]
+
+ if self.skip_prompt and self.next_tokens_are_prompt:
+ self.next_tokens_are_prompt = False
+ return
+
+ # Add the new token to the cache and decodes the entire thing.
+ self.token_cache.extend(value.tolist())
+ text = self.tokenizer.decode(self.token_cache, **self.decode_kwargs)
+
+ # After the symbol for a new line, we flush the cache.
+ if text.endswith("\n"):
+ printable_text = text[self.print_len :]
+ self.token_cache = []
+ self.print_len = 0
+ # If the last token is a CJK character, we print the characters.
+ elif len(text) > 0 and self._is_chinese_char(ord(text[-1])):
+ printable_text = text[self.print_len :]
+ self.print_len += len(printable_text)
+ # Otherwise, prints until the last space char (simple heuristic to avoid printing incomplete words,
+ # which may change with the subsequent token -- there are probably smarter ways to do this!)
+ else:
+ printable_text = text[self.print_len : text.rfind(" ") + 1]
+ self.print_len += len(printable_text)
+
+ self.on_finalized_text(printable_text)
+
+ def end(self):
+ """Flushes any remaining cache and prints a newline to stdout."""
+ # Flush the cache, if it exists
+ if len(self.token_cache) > 0:
+ text = self.tokenizer.decode(self.token_cache, **self.decode_kwargs)
+ printable_text = text[self.print_len :]
+ self.token_cache = []
+ self.print_len = 0
+ else:
+ printable_text = ""
+
+ self.next_tokens_are_prompt = True
+ self.on_finalized_text(printable_text, stream_end=True)
+
+ def on_finalized_text(self, text: str, stream_end: bool = False):
+ """Prints the new text to stdout. If the stream is ending, also prints a newline."""
+ print(text, flush=True, end="" if not stream_end else None)
+
+ def _is_chinese_char(self, cp):
+ """Checks whether CP is the codepoint of a CJK character."""
+ # This defines a "chinese character" as anything in the CJK Unicode block:
+ # https://en.wikipedia.org/wiki/CJK_Unified_Ideographs_(Unicode_block)
+ #
+ # Note that the CJK Unicode block is NOT all Japanese and Korean characters,
+ # despite its name. The modern Korean Hangul alphabet is a different block,
+ # as is Japanese Hiragana and Katakana. Those alphabets are used to write
+ # space-separated words, so they are not treated specially and handled
+ # like the all of the other languages.
+ if (
+ (cp >= 0x4E00 and cp <= 0x9FFF)
+ or (cp >= 0x3400 and cp <= 0x4DBF) #
+ or (cp >= 0x20000 and cp <= 0x2A6DF) #
+ or (cp >= 0x2A700 and cp <= 0x2B73F) #
+ or (cp >= 0x2B740 and cp <= 0x2B81F) #
+ or (cp >= 0x2B820 and cp <= 0x2CEAF) #
+ or (cp >= 0xF900 and cp <= 0xFAFF)
+ or (cp >= 0x2F800 and cp <= 0x2FA1F) #
+ ): #
+ return True
+
+ return False
+
+
+class TextIteratorStreamer(TextStreamer):
+ """
+ Streamer that stores print-ready text in a queue, to be used by a downstream application as an iterator. This is
+ useful for applications that benefit from acessing the generated text in a non-blocking way (e.g. in an interactive
+ Gradio demo).
+
+
+
+ The API for the streamer classes is still under development and may change in the future.
+
+
+
+ Parameters:
+ tokenizer (`AutoTokenizer`):
+ The tokenized used to decode the tokens.
+ skip_prompt (`bool`, *optional*, defaults to `False`):
+ Whether to skip the prompt to `.generate()` or not. Useful e.g. for chatbots.
+ timeout (`float`, *optional*):
+ The timeout for the text queue. If `None`, the queue will block indefinitely. Useful to handle exceptions
+ in `.generate()`, when it is called in a separate thread.
+ decode_kwargs (`dict`, *optional*):
+ Additional keyword arguments to pass to the tokenizer's `decode` method.
+
+ Examples:
+
+ ```python
+ >>> from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+ >>> from threading import Thread
+
+ >>> tok = AutoTokenizer.from_pretrained("openai-community/gpt2")
+ >>> model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2")
+ >>> inputs = tok(["An increasing sequence: one,"], return_tensors="pt")
+ >>> streamer = TextIteratorStreamer(tok)
+
+ >>> # Run the generation in a separate thread, so that we can fetch the generated text in a non-blocking way.
+ >>> generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=20)
+ >>> thread = Thread(target=model.generate, kwargs=generation_kwargs)
+ >>> thread.start()
+ >>> generated_text = ""
+ >>> for new_text in streamer:
+ ... generated_text += new_text
+ >>> generated_text
+ 'An increasing sequence: one, two, three, four, five, six, seven, eight, nine, ten, eleven,'
+ ```
+ """
+
+ def __init__(
+ self, tokenizer: "AutoTokenizer", skip_prompt: bool = False, timeout: Optional[float] = None, **decode_kwargs
+ ):
+ super().__init__(tokenizer, skip_prompt, **decode_kwargs)
+ self.text_queue = Queue()
+ self.stop_signal = None
+ self.timeout = timeout
+
+ def on_finalized_text(self, text: str, stream_end: bool = False):
+ """Put the new text in the queue. If the stream is ending, also put a stop signal in the queue."""
+ self.text_queue.put(text, timeout=self.timeout)
+ if stream_end:
+ self.text_queue.put(self.stop_signal, timeout=self.timeout)
+
+ def __iter__(self):
+ return self
+
+ def __next__(self):
+ value = self.text_queue.get(timeout=self.timeout)
+ if value == self.stop_signal:
+ raise StopIteration()
+ else:
+ return value
+
+
+class AsyncTextIteratorStreamer(TextStreamer):
+ """
+ Streamer that stores print-ready text in a queue, to be used by a downstream application as an async iterator.
+ This is useful for applications that benefit from acessing the generated text asynchronously (e.g. in an
+ interactive Gradio demo).
+
+
+
+ The API for the streamer classes is still under development and may change in the future.
+
+
+
+ Parameters:
+ tokenizer (`AutoTokenizer`):
+ The tokenized used to decode the tokens.
+ skip_prompt (`bool`, *optional*, defaults to `False`):
+ Whether to skip the prompt to `.generate()` or not. Useful e.g. for chatbots.
+ timeout (`float`, *optional*):
+ The timeout for the text queue. If `None`, the queue will block indefinitely. Useful to handle exceptions
+ in `.generate()`, when it is called in a separate thread.
+ decode_kwargs (`dict`, *optional*):
+ Additional keyword arguments to pass to the tokenizer's `decode` method.
+
+ Raises:
+ TimeoutError: If token generation time exceeds timeout value.
+
+ Examples:
+
+ ```python
+ >>> from transformers import AutoModelForCausalLM, AutoTokenizer, AsyncTextIteratorStreamer
+ >>> from threading import Thread
+ >>> import asyncio
+
+ >>> tok = AutoTokenizer.from_pretrained("openai-community/gpt2")
+ >>> model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2")
+ >>> inputs = tok(["An increasing sequence: one,"], return_tensors="pt")
+
+ >>> # Run the generation in a separate thread, so that we can fetch the generated text in a non-blocking way.
+ >>> async def main():
+ ... # Important: AsyncTextIteratorStreamer must be initialized inside a coroutine!
+ ... streamer = AsyncTextIteratorStreamer(tok)
+ ... generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=20)
+ ... thread = Thread(target=model.generate, kwargs=generation_kwargs)
+ ... thread.start()
+ ... generated_text = ""
+ ... async for new_text in streamer:
+ ... generated_text += new_text
+ >>> print(generated_text)
+ >>> asyncio.run(main())
+ An increasing sequence: one, two, three, four, five, six, seven, eight, nine, ten, eleven,
+ ```
+ """
+
+ def __init__(
+ self, tokenizer: "AutoTokenizer", skip_prompt: bool = False, timeout: Optional[float] = None, **decode_kwargs
+ ):
+ super().__init__(tokenizer, skip_prompt, **decode_kwargs)
+ self.text_queue = asyncio.Queue()
+ self.stop_signal = None
+ self.timeout = timeout
+ self.loop = asyncio.get_running_loop()
+ self.has_asyncio_timeout = hasattr(asyncio, "timeout")
+
+ def on_finalized_text(self, text: str, stream_end: bool = False):
+ """Put the new text in the queue. If the stream is ending, also put a stop signal in the queue."""
+ self.loop.call_soon_threadsafe(self.text_queue.put_nowait, text)
+ if stream_end:
+ self.loop.call_soon_threadsafe(self.text_queue.put_nowait, self.stop_signal)
+
+ def __aiter__(self):
+ return self
+
+ async def __anext__(self):
+ try:
+ if self.has_asyncio_timeout:
+ async with asyncio.timeout(self.timeout):
+ value = await self.text_queue.get()
+ else:
+ value = await asyncio.wait_for(self.text_queue.get(), timeout=self.timeout)
+ except asyncio.TimeoutError:
+ raise TimeoutError()
+ else:
+ if value == self.stop_signal:
+ raise StopAsyncIteration()
+ else:
+ return value
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/tensor_parallel.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/tensor_parallel.py
new file mode 100644
index 0000000000000000000000000000000000000000..76215e4666ec67f43e9230eeb965223547bec859
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/tensor_parallel.py
@@ -0,0 +1,672 @@
+# Copyright 2024 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.
+from __future__ import annotations
+
+import re
+from functools import lru_cache, partial
+from typing import List, Optional, Tuple, Union
+
+import torch
+from torch import nn
+
+from ..utils import is_torch_greater_or_equal, logging
+
+
+ALL_LAYERNORM_LAYERS = [nn.LayerNorm]
+
+logger = logging.get_logger(__name__)
+
+# Cache this result has it's a C FFI call which can be pretty time-consuming
+_torch_distributed_available = torch.distributed.is_available()
+
+
+if is_torch_greater_or_equal("2.5") and _torch_distributed_available:
+ from torch.distributed.tensor import DTensor, Placement, Replicate, Shard
+
+
+def _blocks_to_block_sizes(total_size: int, blocks: Union[int, List[int]]) -> List[int]:
+ """
+ Convert block count or proportions to block sizes.
+
+ This function accepts
+
+ - The number of blocks (int), in which case the block size is
+ total_size//blocks; or
+ - A list of block sizes (List[int]).
+
+ In the second case, if sum(blocks) < total_size, the ratios between
+ the block sizes will be preserved. For instance, if blocks is
+ [2, 1, 1] and total_size is 1024, the returned block sizes are
+ [512, 256, 256].
+ """
+ if isinstance(blocks, list):
+ total_blocks = sum(blocks)
+ assert total_size % total_blocks == 0, f"Cannot split {total_size} in proportional blocks: {blocks}"
+ part_size = total_size // total_blocks
+ return [part_size * block for block in blocks]
+ else:
+ assert total_size % blocks == 0, f"Prepacked is not divisible by {blocks}"
+ single_size = total_size // blocks
+ return [single_size] * blocks
+
+
+str_to_torch_dtype = {
+ "BOOL": torch.bool,
+ "U8": torch.uint8,
+ "I8": torch.int8,
+ "I16": torch.int16,
+ "F16": torch.float16,
+ "BF16": torch.bfloat16,
+ "I32": torch.int32,
+ "F32": torch.float32,
+ "F64": torch.float64,
+ "I64": torch.int64,
+ "F8_E4M3": torch.float8_e4m3fn,
+}
+
+
+def get_packed_weights(param, empty_param, device_mesh, rank, dim):
+ """
+ When weights are packed (gate_up_proj), we need to make sure each shard gets its correct share.
+ So if you have: gate_proj ( 16, 5120, 8190)
+ and up_proj ( 16, 5120, 8190)
+ packed as gate_up_proj ( 16, 5120, 2 * 8190)
+ And you shard along the last dimension, you need to interleave the gate and up values:
+
+ Now, if we shard along the last dimension across TP_size (Tensor Parallelism size), we must interleave the values from gate and up projections correctly.
+
+ Let's take TP_size = 4 for an example:
+
+ Packed tensor `gate_up_proj`
+ ---------------------------------------------------------------
+ [ G0 G1 G2 G3 | G4 G5 G6 G7 | ... | U0 U1 U2 U3 | U4 U5 U6 U7 | ... ]
+ ↑─────────────↑ ↑─────────────↑ ↑─────────────↑ ↑─────────────↑
+ Gate Slice 0 Gate Slice 1 Up Slice 0 Up Slice 1
+
+ Explanation:
+ - The first half of the tensor (left of the center) holds the gate_proj values.
+ - The second half (right of the center) holds the up_proj values.
+ - For TP=4, we divide each half into 4 slices. In this example, we show two slices for brevity.
+ - Each shard receives one slice from the gate part and the corresponding slice from the up part.
+
+ For instance:
+ • Shard 0 gets: [ Gate Slice 0, Up Slice 0 ] = [ G0, G1, G2, G3, U0, U1, U2, U3 ]
+ • Shard 1 gets: [ Gate Slice 1, Up Slice 1 ] = [ G4, G5, G6, G7, U4, U5, U6, U7 ]
+ • … and so on.
+
+ This ensures that each shard receives an equal portion of both gate and up projections, maintaining consistency across tensor parallelism.
+ """
+ slice_ = param
+ total_size = empty_param.shape[dim]
+ world_size = device_mesh.size()
+ block_sizes = _blocks_to_block_sizes(total_size=total_size, blocks=2)
+
+ tensors_slices = []
+ block_offset = 0
+ for block_size in block_sizes:
+ shard_block_size = block_size // world_size
+ start = rank * shard_block_size
+ stop = (rank + 1) * shard_block_size
+ tensors_slices += range(block_offset + start, block_offset + stop)
+ block_offset += block_size
+
+ slice_dtype = slice_.get_dtype()
+ # Handle F8_E4M3 dtype by converting to float16 before slicing
+ # Without upcasting, the slicing causes : RuntimeError: "index_cpu" not implemented for 'Float8_e4m3fn'
+ if slice_dtype == "F8_E4M3":
+ slice_ = slice_[...].to(torch.float16)
+
+ if dim == 0:
+ tensor = slice_[tensors_slices, ...]
+ elif dim == 1 or dim == -2:
+ tensor = slice_[:, tensors_slices, ...]
+ elif dim == 2 or dim == -1:
+ tensor = slice_[..., tensors_slices]
+ else:
+ raise ValueError(f"Unsupported dim {dim}, only dim 0, 1 or 2 are supported")
+ return tensor.to(str_to_torch_dtype[slice_dtype])
+
+
+def get_tensor_shard(param, empty_param, device_mesh, rank, dim):
+ if dim == 0:
+ size_ = empty_param.shape[0]
+ param = param[rank * (size_ // device_mesh.size()) : (rank + 1) * (size_ // device_mesh.size()), ...]
+ elif dim == 1 or dim == -2:
+ size_ = empty_param.shape[-2]
+ param = param[..., rank * (size_ // device_mesh.size()) : (rank + 1) * (size_ // device_mesh.size()), :]
+ elif dim == 2 or dim == -1:
+ size_ = empty_param.shape[-1]
+ param = param[..., rank * (size_ // device_mesh.size()) : (rank + 1) * (size_ // device_mesh.size())]
+ else:
+ raise ValueError(f"Unsupported dim {dim}, only dim 0, 1 or 2 are supported")
+ return param
+
+
+def distribute_module(
+ module: nn.Module,
+ device_mesh=None,
+ input_fn=None,
+ output_fn=None,
+) -> nn.Module:
+ """
+ Copy pasted from torch's function but we remove the communications (partitionning)
+ as well as buffer registering that is similarly not efficient.
+ """
+ if len(module._forward_pre_hooks) == 0:
+ if input_fn is not None:
+ module.register_forward_pre_hook(lambda mod, inputs: input_fn(mod, inputs, device_mesh))
+ if output_fn is not None:
+ module.register_forward_hook(lambda mod, inputs, outputs: output_fn(mod, outputs, device_mesh))
+ return module
+
+
+class TensorParallelLayer:
+ """
+ General tensor parallel layer for transformers.
+ """
+
+ use_dtensor = True
+
+ @staticmethod
+ def _prepare_input_fn(input_layouts, desired_input_layouts, mod, inputs, device_mesh): ...
+
+ @staticmethod
+ def _prepare_output_fn(output_layouts, use_local_output, mod, outputs, device_mesh): ...
+
+ def partition_tensor(self, param, empty_param, param_type, param_casting_dtype, to_contiguous, rank, device_mesh):
+ raise NotImplementedError
+
+ def prepare_module_tp(self, module: nn.Module, device_mesh) -> nn.Module:
+ if self.use_dtensor:
+ distribute_module(
+ module,
+ device_mesh,
+ partial(self._prepare_input_fn, self.input_layouts, self.desired_input_layouts),
+ partial(self._prepare_output_fn, self.output_layouts, self.use_local_output),
+ )
+
+
+# use_dtensor needs to be set to false for nn.Parameter when you want to view, chunk, slice
+# you name it. Whatever you want to do that is a bit unconventional, you need local tensors
+class GatherParallel(TensorParallelLayer):
+ """
+ Simple class used to define the hooks to add to a layer when we just want to gather the outputs
+ """
+
+ def __init__(
+ self,
+ *,
+ input_layouts: Optional[Placement] = None,
+ output_layouts: Optional[Placement] = None,
+ use_local_output: bool = True,
+ ):
+ super().__init__()
+ self.input_layouts = (input_layouts or Replicate(),)
+ self.output_layouts = output_layouts
+ self.desired_input_layouts = (Replicate(),)
+ self.use_local_output = use_local_output
+
+ @staticmethod
+ def _prepare_input_fn(input_layouts, desired_input_layouts, mod, inputs, device_mesh):
+ if inputs and isinstance(inputs[0], DTensor):
+ inputs = inputs[0].to_local()
+ return inputs
+
+ @staticmethod
+ def _prepare_output_fn(output_layouts, use_local_output, mod, outputs, device_mesh):
+ # this op cannot be asynch, otherwise it completely breaks the outputs of models
+ torch.distributed.all_reduce(outputs[0], op=torch.distributed.ReduceOp.SUM, async_op=False)
+ return outputs
+
+
+class IsolatedParallel(TensorParallelLayer):
+ """
+ This class is used to isolate computation in a TP layer from the rest of the world.
+ Parameters need to be LOCAL, so not dtensors
+ """
+
+ @staticmethod
+ def _prepare_input_fn(input_layouts, desired_input_layouts, mod, inputs, device_mesh=None):
+ # annotate module input placements/sharding with input_layouts
+ input_tensor = inputs[0]
+ if isinstance(input_tensor, DTensor):
+ input_tensor = input_tensor.to_local()
+ return input_tensor
+
+ @staticmethod
+ def _prepare_output_fn(output_layouts, use_local_output, mod, outputs, device_mesh=None):
+ # TODO: figure out dynamo support for instance method and switch this to instance method
+ return outputs
+
+ def prepare_module_tp(self, module: nn.Module, device_mesh) -> nn.Module:
+ distribute_module(
+ module,
+ device_mesh,
+ partial(self._prepare_input_fn, None, None),
+ partial(self._prepare_output_fn, None, None),
+ )
+
+
+class ColwiseParallel(TensorParallelLayer):
+ """
+ General tensor parallel layer for transformers.
+ """
+
+ def __init__(
+ self,
+ *,
+ input_layouts: Optional[Placement] = None,
+ output_layouts: Optional[Placement] = None,
+ use_local_output: bool = True,
+ use_dtensor=True,
+ ):
+ super().__init__()
+ self.input_layouts = (input_layouts or Replicate(),)
+ self.output_layouts = (output_layouts or Shard(-1),)
+ self.desired_input_layouts = (Replicate(),)
+ self.use_local_output = use_local_output
+ self.use_dtensor = use_dtensor
+
+ @staticmethod
+ def _prepare_input_fn(input_layouts, desired_input_layouts, mod, inputs, device_mesh):
+ # TODO: figure out dynamo support for instance method and switch this to instance method
+ # annotate module input placements/sharding with input_layouts
+ input_tensor = inputs[0]
+ if not isinstance(input_tensor, DTensor):
+ input_tensor = DTensor.from_local(input_tensor, device_mesh, input_layouts, run_check=False)
+
+ # transform the input layouts to the desired layouts of ColwiseParallel
+ if input_layouts != desired_input_layouts:
+ input_tensor = input_tensor.redistribute(placements=desired_input_layouts, async_op=False)
+ return input_tensor
+
+ def partition_tensor(self, param, empty_param, param_type, param_casting_dtype, to_contiguous, rank, device_mesh):
+ # colwise shard weight/bias to Shard(0), weight be Shard(-2) (0 if you have 1 dim only)
+ # means Colwise as Linear is input * weight^T + bias, where
+ # weight would become Shard(1)
+ if param_type == "bias":
+ parameter = get_tensor_shard(param, empty_param, device_mesh, rank, -1)
+ shard = [Shard(-1)]
+ else:
+ shard = [Shard(-2)]
+ parameter = get_tensor_shard(param, empty_param, device_mesh, rank, -2)
+
+ parameter = parameter.to(param_casting_dtype)
+ if to_contiguous:
+ parameter = parameter.contiguous()
+ if self.use_dtensor:
+ parameter = DTensor.from_local(parameter, device_mesh, shard, run_check=False)
+ return nn.Parameter(parameter)
+
+ @staticmethod
+ def _prepare_output_fn(output_layouts, use_local_output, mod, outputs, device_mesh):
+ # outputs is a shard on last dimension DTensor, i.e. Shard(-1)
+ if outputs.placements != output_layouts:
+ outputs = outputs.redistribute(placements=output_layouts, async_op=False)
+ # back to local tensor
+ return outputs.to_local() if use_local_output else outputs
+
+
+class PackedColwiseParallel(ColwiseParallel):
+ def partition_tensor(self, param, empty_param, param_type, param_casting_dtype, to_contiguous, rank, device_mesh):
+ # colwise shard weight/bias to Shard(0), weight be Shard(-2) (0 if you have 1 dim only)
+ # means Colwise as Linear is input * weight^T + bias, where
+ # weight would become Shard(1)
+ parameter = get_packed_weights(param, empty_param, device_mesh, rank, -2)
+ parameter = parameter.to(param_casting_dtype)
+ if to_contiguous:
+ parameter = parameter.contiguous()
+ if self.use_dtensor:
+ parameter = DTensor.from_local(parameter, device_mesh, [Shard(-2)], run_check=False)
+ return nn.Parameter(parameter)
+
+
+class RowwiseParallel(TensorParallelLayer):
+ """
+ Partition a compatible nn.Module in a row-wise fashion. Currently supports nn.Linear and nn.Embedding.
+ Users can compose it with ColwiseParallel to achieve the sharding of more complicated modules.
+ (i.e. MLP, Attention)
+
+ Keyword Args:
+ input_layouts (Placement, optional):
+ The DTensor layout of input tensor for the nn.Module, this is used to annotate the input tensor to
+ become a DTensor. If not specified, we assume the input tensor to be sharded on the last dimension.
+ output_layouts (Placement, optional):
+ The DTensor layout of the output for the nn.Module, this is used to ensure the output of the nn.Module
+ with the user desired layout. If not specified, the output tensor is replicated.
+ use_local_output (bool, optional):
+ Whether to use local :class:`torch.Tensor` instead of :class:`DTensor` for the module output, default: True.
+ Returns:
+ A :class:`ParallelStyle` object that represents Rowwise sharding of the nn.Module.
+ """
+
+ def __init__(
+ self,
+ *,
+ input_layouts: Optional[Placement] = None,
+ output_layouts: Optional[Placement] = None,
+ use_local_output: bool = True,
+ use_dtensor=True,
+ ):
+ super().__init__()
+ self.input_layouts = (input_layouts or Shard(-1),)
+ self.output_layouts = (output_layouts or Replicate(),)
+ self.use_local_output = use_local_output
+ self.use_dtensor = use_dtensor
+
+ def partition_tensor(self, param, empty_param, param_type, param_casting_dtype, to_contiguous, rank, device_mesh):
+ # Rowwise shard weight to Shard(1), bias to Replicate(), weight be Shard(1)
+ # means Rowwise as nn.Linear is input * weight^T + bias, where
+ # weight would become Shard(0)
+ if param_type != "bias":
+ parameter = get_tensor_shard(param, empty_param, device_mesh, rank, -1)
+ shard = [Shard(-1)]
+ else:
+ shard = [Replicate()]
+ parameter = param[:]
+
+ parameter = parameter.to(param_casting_dtype)
+ if to_contiguous:
+ parameter = parameter.contiguous()
+ if self.use_dtensor:
+ parameter = DTensor.from_local(parameter, device_mesh, shard, run_check=False)
+ return nn.Parameter(parameter)
+
+ @staticmethod
+ def _prepare_input_fn(input_layouts, desired_input_layouts, mod, inputs, device_mesh):
+ if hasattr(mod, "bias") and mod.bias is not None:
+ mod._bias = mod.bias
+ mod.bias = None
+
+ input_tensor = inputs[0]
+ if not isinstance(input_tensor, DTensor):
+ input_tensor = DTensor.from_local(input_tensor, device_mesh, input_layouts, run_check=False)
+
+ if input_layouts != desired_input_layouts:
+ input_tensor = input_tensor.redistribute(placements=desired_input_layouts, async_op=True)
+ return input_tensor
+
+ @staticmethod
+ def _prepare_output_fn(output_layouts, use_local_output, mod, outputs, device_mesh):
+ # Rowwise sharding produces partial output, depending on output layouts:
+ # 1. to replicate -> allreduce
+ # 2. to shard -> reduce_scatter
+ if outputs.placements != output_layouts:
+ outputs = outputs.redistribute(placements=output_layouts, async_op=True)
+ if hasattr(mod, "_bias"):
+ outputs += mod._bias
+ # back to local tensor if use_local_output is True
+ return outputs.to_local() if use_local_output else outputs
+
+ def prepare_module_tp(self, module: nn.Module, device_mesh) -> nn.Module:
+ module._distribute_module_applied = True
+ if self.use_dtensor:
+ if isinstance(module, nn.Linear):
+ # rowwise linear runtime sharding requires input tensor shard on last dim
+ self.desired_input_layouts: Tuple[Placement, ...] = (Shard(-1),)
+ elif isinstance(module, nn.Embedding):
+ # rowwise embedding runtime sharding requires input tensor replicated
+ self.desired_input_layouts = (Replicate(),)
+ elif isinstance(module, nn.Parameter):
+ # rowwise embedding runtime sharding requires input tensor replicated
+ self.desired_input_layouts = (Shard(-1),)
+ else:
+ raise NotImplementedError("RowwiseParallel currently only support nn.Linear and nn.Embedding!")
+
+ distribute_module(
+ module,
+ device_mesh,
+ partial(self._prepare_input_fn, self.input_layouts, self.desired_input_layouts),
+ partial(self._prepare_output_fn, self.output_layouts, self.use_local_output),
+ )
+
+
+class PackedRowwiseParallel(RowwiseParallel):
+ def partition_tensor(self, param, empty_param, param_type, param_casting_dtype, to_contiguous, rank, device_mesh):
+ # colwise shard weight/bias to Shard(0), weight be Shard(-2) (0 if you have 1 dim only)
+ # means Colwise as Linear is input * weight^T + bias, where
+ # weight would become Shard(1)
+ parameter = get_packed_weights(param, empty_param, device_mesh, rank, -1)
+ parameter = parameter.to(param_casting_dtype)
+ if to_contiguous:
+ parameter = parameter.contiguous()
+ if self.use_dtensor:
+ parameter = DTensor.from_local(parameter, device_mesh, [Shard(-1)], run_check=False)
+ return nn.Parameter(parameter)
+
+
+class SequenceParallel(TensorParallelLayer):
+ """
+ SequenceParallel replicates a compatible ``nn.Module`` parameters and runs the sharded computation with
+ input sharded on the sequence dimension. This currently supports ``nn.LayerNorm``, ``nn.Dropout``, and the
+ `RMSNorm python implementation `__
+
+ This style implements the operation that is described in the paper
+ `Reducing Activation Recomputation in Large Transformer Models `__
+
+ If the input passed in to this ``nn.Module`` is a :class:`torch.Tensor`, it assumes that the input is already sharded
+ on the sequence dimension and converts the input to a :class:`DTensor` sharded on the sequence dimension. If the input
+ passed in to this ``nn.Module`` is already a :class:`DTensor` but is not sharded on the sequence dimension, it would
+ redistribute the input to be sharded on the sequence dimension.
+
+ The output of the ``nn.Module`` will be sharded on the sequence dimension.
+
+ Keyword Args:
+ sequence_dim (int, optional):
+ The sequence dimension of the input tensor for the ``nn.Module``, this is used to annotate the input tensor to
+ become a DTensor that is sharded on the sequence dimension, default: 1.
+ use_local_output (bool, optional):
+ Whether to use local :class:`torch.Tensor` instead of :class:`DTensor` for the module output, default: False.
+ Returns:
+ A :class:`ParallelStyle` object that represents Sequence Parallel of the ``nn.Module``.
+
+ Example::
+ >>> # xdoctest: +SKIP(failing)
+ >>> from torch.distributed.tensor.parallel import parallelize_module, SequenceParallel
+ >>> from torch.distributed.device_mesh import init_device_mesh
+ >>> ...
+ >>> m = Model(...) # m is a nn.Module that contains a "norm" nn.LayerNorm submodule
+ >>> tp_mesh = init_device_mesh("cuda", (8,))
+ >>>
+ >>> # By default, the input of the "norm" will be converted to DTensor that shards on the sequence dim
+ >>> # and the output of "norm" will return a sharded on sequence dimension :class:`DTensor`.
+ >>>
+ >>> sharded_mod = parallelize_module(m, tp_mesh, {"norm": SequenceParallel()}),
+ >>> ...
+
+ .. note:: SequenceParallel style assumes ones initialization if there are weights in the nn.Module (i.e.
+ ``nn.LayerNorm`` or ``RMSNorm``, and they by default have ones initialization). If you have custom
+ inits for the weights on those modules, you need to broadcast the weights before/after parallelizing
+ to ensure that they are replicated.
+ """
+
+ def __init__(self, *, sequence_dim: int = 1, use_local_output: bool = False, use_dtensor=False):
+ super().__init__()
+ self.input_layouts = (Replicate(),)
+ self.desired_input_layouts = (Shard(1),)
+ self.output_layouts = (Replicate(),)
+ self.use_local_output = use_local_output
+ self.use_dtensor = True
+ self.sequence_sharding = (Shard(sequence_dim),)
+ self.use_local_output = use_local_output
+
+ @staticmethod
+ def _prepare_input_fn(input_layouts, desired_input_layouts, mod, inputs, device_mesh):
+ input_tensor = inputs[0]
+ if not isinstance(input_tensor, DTensor):
+ input_tensor = DTensor.from_local(input_tensor, device_mesh, input_layouts, run_check=False)
+ if input_layouts != desired_input_layouts:
+ input_tensor = input_tensor.redistribute(placements=desired_input_layouts, async_op=True)
+ return input_tensor
+
+ @staticmethod
+ def _prepare_output_fn(output_layouts, use_local_output, mod, outputs, device_mesh):
+ outputs = outputs.redistribute(
+ placements=(Replicate(),), async_op=True
+ ) # maybe we have to replicate ? because next layer is not sharded
+ return outputs.to_local() # if use_local_output else outputs
+
+ def partition_tensor(self, param, empty_param, param_type, param_casting_dtype, to_contiguous, rank, device_mesh):
+ # colwise shard weight/bias to Shard(0), weight be Shard(-2) (0 if you have 1 dim only)
+ # means Colwise as Linear is input * weight^T + bias, where
+ # weight would become Shard(1)
+ parameter = param[:]
+ parameter = parameter.to(param_casting_dtype)
+ if to_contiguous:
+ parameter = parameter.contiguous()
+ if self.use_dtensor:
+ parameter = DTensor.from_local(parameter, device_mesh, [Replicate()], run_check=False)
+ return nn.Parameter(parameter)
+
+
+SUPPORTED_TP_STYLES = {
+ "colwise",
+ "rowwise",
+ "colwise_rep",
+ "rowwise_rep",
+ "local_colwise",
+ "local_rowwise",
+ "local",
+ "gather",
+ "local_packed_rowwise",
+ "sequence_parallel",
+}
+
+
+@lru_cache
+def translate_to_torch_parallel_style(style: str):
+ """
+ In model configurations, we use a neutral type (string) to specify parallel
+ styles, here we translate them into torch.distributed tensor-parallel
+ types.
+ """
+ if not isinstance(style, str):
+ raise ValueError(f"Unsupported parallel style type {type(style)}, expected str")
+
+ if style == "colwise":
+ return ColwiseParallel()
+ elif style == "rowwise":
+ return RowwiseParallel()
+ elif style == "colwise_rep":
+ return ColwiseParallel(output_layouts=Replicate())
+ elif style == "rowwise_rep":
+ return RowwiseParallel(input_layouts=Replicate())
+ elif style == "local_colwise":
+ return ColwiseParallel(use_dtensor=False)
+ elif style == "local_rowwise":
+ return RowwiseParallel(use_dtensor=False)
+ elif style == "local":
+ return IsolatedParallel()
+ elif style == "gather":
+ return GatherParallel()
+ elif style == "local_packed_rowwise":
+ return PackedRowwiseParallel(use_dtensor=False)
+ elif style == "sequence_parallel":
+ return SequenceParallel()
+ else:
+ raise ValueError(f"Unsupported parallel style value: {style}")
+
+
+def add_tensor_parallel_hooks_to_module(model, module, tp_plan, layer_name, current_module_plan, device_mesh):
+ """
+ Add hooks to the module holding the layer. Meaning:
+ ```
+ class MyModel(nn.Module):
+ def __init__(self):
+ self.layer = nn.Linear(10, 10)
+ ```
+ has state_dict like:
+ ```
+ {
+ "layer.weight": torch.Tensor,
+ "layer.bias": torch.Tensor
+ }
+ ```
+ we add hooks to `MyModel` as well as `layer` to make sure that the tensors are correctly sharded and gathered.
+ """
+
+ # 1. We add hooks to the layer being loaded:
+ if current_module_plan is not None:
+ tp_layer = translate_to_torch_parallel_style(current_module_plan)
+ try:
+ tp_layer.prepare_module_tp(module, device_mesh)
+ except NotImplementedError as e:
+ print(
+ f"Trying to prepare {layer_name}, but it's not supported. Corresponding module: {module} Fix it's TP plan: {e}"
+ )
+
+ # 2. We add hooks to the parrent module if needed
+ if "." in layer_name:
+ parrent_layer_name = layer_name.rsplit(".", 1)[0]
+ generic_name = re.sub(r"\d+", "*", parrent_layer_name)
+ # The module itself needs hooks
+ if module_plan := tp_plan.get(generic_name, False):
+ tp_layer = translate_to_torch_parallel_style(module_plan)
+ module_to_tp_ = model.get_submodule(parrent_layer_name)
+ tp_layer.prepare_module_tp(module_to_tp_, device_mesh)
+
+
+def shard_and_distribute_module(
+ model, param, empty_param, parameter_name, param_casting_dtype, is_contiguous, rank, device_mesh
+):
+ r"""
+ Main uses cases:
+ - column / rowise parallelism, you just shard all the weights of the layer (weight and bias)
+ - packed layers: you slice the weights, then shard like above
+ - custom operation:
+ - you want to add an all-gather at the end of a local layer.
+ - you want to have a layer that is isolated from the rest of the world (because torch.DTensor does not work well with `.view` for instance)
+
+ """
+ param_name, param_type = parameter_name.rsplit(".", 1) if "." in parameter_name else parameter_name
+ tp_plan = model._tp_plan
+ module_to_tp = model.get_submodule(param_name)
+ current_module_plan = None
+ rank = int(rank)
+ generic_param_name = re.sub(r"\d+", "*", parameter_name)
+ if generic_param_name in tp_plan:
+ current_module_plan = tp_plan[generic_param_name]
+ elif "." in generic_param_name and generic_param_name.rsplit(".", 1)[0] in tp_plan:
+ current_module_plan = tp_plan[generic_param_name.rsplit(".", 1)[0]]
+
+ # Add hooks to the module if not done yet
+ # add_tensor_parallel_hooks_to_module(model, module_to_tp, tp_plan, param_name, current_module_plan, device_mesh)
+ if not getattr(module_to_tp, "_is_hooked", False):
+ add_tensor_parallel_hooks_to_module(model, module_to_tp, tp_plan, param_name, current_module_plan, device_mesh)
+ module_to_tp._is_hooked = True
+
+ if current_module_plan is not None:
+ try:
+ tp_layer = translate_to_torch_parallel_style(current_module_plan)
+ param = tp_layer.partition_tensor(
+ param, empty_param, param_type, param_casting_dtype, is_contiguous, rank, device_mesh
+ )
+ except NotImplementedError as e:
+ print(
+ f"Trying to prepare {parameter_name}, but it's not supported. Corresponding module: {module_to_tp} Fix it's TP plan, current layer: {tp_layer} : {e}"
+ )
+ else:
+ # TODO log no plan modules in set
+ # print("No plan for", parameter_name,end ="\n")
+ param = param[...].to(param_casting_dtype)
+ if is_contiguous:
+ param = param.contiguous()
+
+ # SUPER IMPORTANT we have to use setattr
+ # otherwise loading is crazy slow
+ if not isinstance(param, torch.nn.Parameter):
+ param = torch.nn.Parameter(param)
+ setattr(module_to_tp, param_type, param)
+ # module_to_tp.load_state_dict({param_type: param}, strict=False, assign=True)
+ return param
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/tf_logits_process.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/tf_logits_process.py
new file mode 100644
index 0000000000000000000000000000000000000000..f70655fb7c1d0b502dfc8f1c4f2c491867aa7e8c
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/tf_logits_process.py
@@ -0,0 +1,603 @@
+# coding=utf-8
+# Copyright 2022 The HuggingFace Inc. 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.
+
+import inspect
+from typing import List, Tuple
+
+import numpy as np
+import tensorflow as tf
+
+from ..tf_utils import stable_softmax
+from ..utils import add_start_docstrings
+from ..utils.logging import get_logger
+
+
+logger = get_logger(__name__)
+
+
+TF_LOGITS_PROCESSOR_INPUTS_DOCSTRING = r"""
+ Args:
+ input_ids (`tf.Tensor` of shape `(batch_size, sequence_length)`):
+ Indices of input sequence tokens in the vocabulary.
+
+ Indices can be obtained using [`PreTrainedTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+ [`PreTrainedTokenizer.__call__`] for details.
+
+ [What are input IDs?](../glossary#input-ids)
+ scores (`tf.Tensor` of shape `(batch_size, config.vocab_size)`):
+ Prediction scores of a language modeling head. These can be logits for each vocabulary when not using beam
+ search or log softmax for each vocabulary token when using beam search.
+ cur_len (`int`):
+ The current length of valid input sequence tokens. In the TF implementation, the input_ids' sequence length
+ is the maximum length generate can produce, and we need to know which of its tokens are valid.
+ kwargs (`Dict[str, Any]`, *optional*):
+ Additional logits processor specific kwargs.
+
+ Return:
+ `tf.Tensor` of shape `(batch_size, config.vocab_size)`: The processed prediction scores.
+"""
+
+
+class TFLogitsProcessor:
+ """Abstract base class for all logit processors that can be applied during generation."""
+
+ @add_start_docstrings(TF_LOGITS_PROCESSOR_INPUTS_DOCSTRING)
+ def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor:
+ """TF method for processing logits."""
+ raise NotImplementedError(
+ f"{self.__class__} is an abstract class. Only classes inheriting this class can be called."
+ )
+
+
+class TFLogitsWarper:
+ """Abstract base class for all logit warpers that can be applied during generation with multinomial sampling."""
+
+ @add_start_docstrings(TF_LOGITS_PROCESSOR_INPUTS_DOCSTRING)
+ def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor:
+ """TF method for warping logits."""
+ raise NotImplementedError(
+ f"{self.__class__} is an abstract class. Only classes inheriting this class can be called."
+ )
+
+
+class TFLogitsProcessorList(list):
+ """
+ This class can be used to create a list of [`TFLogitsProcessor`] to subsequently process a `scores` input tensor.
+ This class inherits from list and adds a specific *__call__* method to apply each [`TFLogitsProcessor`] to the
+ inputs.
+ """
+
+ @add_start_docstrings(TF_LOGITS_PROCESSOR_INPUTS_DOCSTRING)
+ def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int, **kwargs) -> tf.Tensor:
+ for processor in self:
+ function_args = inspect.signature(processor.__call__).parameters
+ if len(function_args) > 3:
+ if not all(arg in kwargs for arg in list(function_args.keys())[2:]):
+ raise ValueError(
+ f"Make sure that all the required parameters: {list(function_args.keys())} for "
+ f"{processor.__class__} are passed to the logits processor."
+ )
+ scores = processor(input_ids, scores, cur_len, **kwargs)
+ else:
+ scores = processor(input_ids, scores, cur_len)
+ return scores
+
+
+class TFTemperatureLogitsWarper(TFLogitsWarper):
+ r"""
+ [`TFLogitsWarper`] for temperature (exponential scaling output probability distribution).
+
+ Args:
+ temperature (`float`):
+ The value used to module the logits distribution.
+ """
+
+ def __init__(self, temperature: float):
+ if not isinstance(temperature, float) or not (temperature > 0):
+ raise ValueError(f"`temperature` has to be a strictly positive float, but is {temperature}")
+
+ self.temperature = temperature
+
+ def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor:
+ scores = scores / self.temperature
+ return scores
+
+
+class TFTopKLogitsWarper(TFLogitsWarper):
+ r"""
+ [`TFLogitsWarper`] that performs top-k, i.e. restricting to the k highest probability elements.
+
+ Args:
+ top_k (`int`):
+ The number of highest probability vocabulary tokens to keep for top-k-filtering.
+ filter_value (`float`, *optional*, defaults to -inf):
+ All filtered values will be set to this float value.
+ min_tokens_to_keep (`int`, *optional*, defaults to 1):
+ Minimum number of tokens that cannot be filtered.
+ """
+
+ def __init__(self, top_k: int, filter_value: float = -float("Inf"), min_tokens_to_keep: int = 1):
+ if not isinstance(top_k, int) or top_k <= 0:
+ raise ValueError(f"`top_k` has to be a strictly positive integer, but is {top_k}")
+
+ self.top_k = max(top_k, min_tokens_to_keep)
+ self.filter_value = filter_value
+
+ def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor:
+ top_k = min(self.top_k, scores.shape[-1]) # Safety check
+ # Boolean mask containing all tokens with a probability less than the last token of the top-k
+ indices_to_remove = scores < tf.math.top_k(scores, k=top_k)[0][..., -1:]
+ next_scores = tf.where(indices_to_remove, self.filter_value, scores)
+ return next_scores
+
+
+class TFTopPLogitsWarper(TFLogitsWarper):
+ """
+ [`TFLogitsWarper`] that performs top-p, i.e. restricting to top tokens summing to <= prob_cut_off.
+
+ Args:
+ top_p (`float`):
+ If set to < 1, only the smallest set of most probable tokens with probabilities that add up to `top_p` or
+ higher are kept for generation.
+ filter_value (`float`, *optional*, defaults to -inf):
+ All filtered values will be set to this float value.
+ min_tokens_to_keep (`int`, *optional*, defaults to 1):
+ Minimum number of tokens that cannot be filtered.
+ """
+
+ def __init__(self, top_p: float, filter_value: float = -float("Inf"), min_tokens_to_keep: int = 1):
+ if not isinstance(top_p, float) or (top_p < 0 or top_p > 1.0):
+ raise ValueError(f"`top_p` has to be a float > 0 and < 1, but is {top_p}")
+ if not isinstance(min_tokens_to_keep, int) or (min_tokens_to_keep < 1):
+ raise ValueError(f"`min_tokens_to_keep` has to be a positive integer, but is {min_tokens_to_keep}")
+
+ self.top_p = top_p
+ self.filter_value = filter_value
+ self.min_tokens_to_keep = min_tokens_to_keep
+
+ def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor:
+ topk_scores, topk_indices = tf.math.top_k(scores, scores.shape[-1])
+
+ mask_scores = tf.fill(scores.shape, self.filter_value)
+ cumulative_probs = tf.math.cumsum(stable_softmax(topk_scores, axis=-1), axis=-1)
+ score_mask = cumulative_probs < self.top_p
+
+ # Also include the token that is higher than top_p (the first false = shift and insert a True on the left)
+ score_mask = tf.concat((tf.ones([score_mask.shape[0], 1], dtype=tf.bool), score_mask[:, :-1]), axis=-1)
+
+ # Ensure min tokens to keep
+ score_mask = tf.concat(
+ (
+ tf.ones([score_mask.shape[0], self.min_tokens_to_keep], dtype=tf.bool),
+ score_mask[:, self.min_tokens_to_keep :],
+ ),
+ axis=-1,
+ )
+
+ # Mask the values that do not fit the criteria
+ topk_next_scores = tf.where(score_mask, topk_scores, mask_scores)
+
+ # Undo the topk sorting: converts the 2D matrix of per-row original indices of shape (batch_size, vocab_size)
+ # to a 3D tensor of shape (batch_size, vocab_size, 2) containing the original score coordinate, from which we
+ # can scatter (i.e. `scatter_indices[row, col, :]` is a tensor containing `[row, topk_indices[row, col]]`)
+ scatter_rows = tf.tile(tf.expand_dims(tf.range(topk_indices.shape[0]), axis=-1), [1, topk_indices.shape[-1]])
+ scatter_indices = tf.stack((scatter_rows, topk_indices), axis=-1)
+ next_scores = tf.scatter_nd(scatter_indices, topk_next_scores, shape=topk_next_scores.shape)
+
+ return next_scores
+
+
+class TFMinLengthLogitsProcessor(TFLogitsProcessor):
+ r"""
+ [`TFLogitsProcessor`] enforcing a min-length by setting EOS probability to 0.
+
+ Args:
+ min_length (`int`):
+ The minimum length below which the score of `eos_token_id` is set to `-float("Inf")`.
+ eos_token_id (`int`):
+ The id of the *end-of-sequence* token.
+ """
+
+ def __init__(self, min_length: int, eos_token_id: int):
+ if not isinstance(min_length, int) or min_length < 0:
+ raise ValueError(f"`min_length` has to be a positive integer, but is {min_length}")
+
+ if not isinstance(eos_token_id, int) or eos_token_id < 0:
+ raise ValueError(f"`eos_token_id` has to be a positive integer, but is {eos_token_id}")
+
+ self.min_length = min_length
+ self.eos_token_id = eos_token_id
+
+ def _apply_eos_token_mask(self, scores: tf.Tensor) -> tf.Tensor:
+ eos_token_id_mask = tf.range(scores.shape[-1]) == self.eos_token_id
+ scores = tf.where(eos_token_id_mask, float("-inf"), scores)
+ return scores
+
+ def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor:
+ # applies eos token masking if the first argument is true
+ scores = tf.cond(
+ tf.less(cur_len, self.min_length),
+ lambda: self._apply_eos_token_mask(scores),
+ lambda: tf.identity(scores),
+ )
+ return scores
+
+
+class TFRepetitionPenaltyLogitsProcessor(TFLogitsProcessor):
+ r"""
+ [`TFLogitsProcessor`] enforcing an exponential penalty on repeated sequences.
+
+ Args:
+ repetition_penalty (`float`):
+ The parameter for repetition penalty. 1.0 means no penalty. See [this
+ paper](https://arxiv.org/pdf/1909.05858.pdf) for more details.
+ """
+
+ def __init__(self, penalty: float):
+ if not isinstance(penalty, float) or not (penalty > 0):
+ raise ValueError(f"`penalty` has to be a strictly positive float, but is {penalty}")
+
+ self.penalty = penalty
+
+ def _create_score_penalties(self, input_ids: tf.Tensor, logits: tf.Tensor) -> tf.Tensor:
+ # We want to populate the penalties in the positions of `input_ids`. Since XLA can't handle shapes unknown
+ # before runtime, `tf.unique` can't be used. Therefore, we may have redundant updates, when a given row has
+ # the same token multiple times.
+
+ # Gathers the penalties to apply
+ logit_penalties = tf.gather(logits, input_ids, axis=1, batch_dims=1)
+ logit_penalties = tf.where(logit_penalties > 0, 1 / self.penalty, logit_penalties)
+ logit_penalties = tf.where(logit_penalties < 0, self.penalty, logit_penalties)
+
+ # Scatters the penalties
+ token_penalties = tf.ones(logits.shape)
+ batch_size = input_ids.shape[0]
+ seq_len = tf.shape(input_ids)[1] # the sequence length has dynamic size, hence the dynamic shape
+ indexable_prev_input_ids = tf.concat(
+ (
+ tf.expand_dims(tf.repeat(tf.range(batch_size), seq_len), axis=-1),
+ tf.expand_dims(tf.reshape(input_ids, [-1]), axis=-1),
+ ),
+ axis=1,
+ )
+ token_penalties = tf.tensor_scatter_nd_update(
+ token_penalties, indices=indexable_prev_input_ids, updates=tf.reshape(logit_penalties, [-1])
+ )
+ return token_penalties
+
+ def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor:
+ score_penalties = self._create_score_penalties(input_ids[:, :cur_len], scores)
+
+ scores = tf.math.multiply(scores, score_penalties)
+
+ return scores
+
+
+class TFNoBadWordsLogitsProcessor(TFLogitsProcessor):
+ """
+ [`TFLogitsProcessor`] that enforces that specified sequences will never be sampled.
+
+ Args:
+ bad_words_ids (`List[List[int]]`):
+ List of list of token ids that are not allowed to be generated. In order to get the tokens of the words
+ that should not appear in the generated text, make sure to set `add_prefix_space=True` when initializing
+ the tokenizer, and use `tokenizer(bad_words, add_special_tokens=False).input_ids`. The `add_prefix_space`
+ argument is only supported for some slow tokenizers, as fast tokenizers' prefixing behaviours come from
+ `pre tokenizers`. Read more [here](https://huggingface.co/docs/tokenizers/api/pre-tokenizers).
+ eos_token_id (`int`):
+ The id of the *end-of-sequence* token.
+ """
+
+ def __init__(self, bad_words_ids: List[List[int]], eos_token_id: int):
+ if not isinstance(bad_words_ids, List) or len(bad_words_ids) == 0:
+ raise ValueError(f"`bad_words_ids` has to be a non-empty list, but is {bad_words_ids}.")
+ if any(not isinstance(bad_word_ids, list) for bad_word_ids in bad_words_ids):
+ raise ValueError(f"`bad_words_ids` has to be a list of lists, but is {bad_words_ids}.")
+ if any(
+ any((not isinstance(token_id, (int, np.integer)) or token_id < 0) for token_id in bad_word_ids)
+ for bad_word_ids in bad_words_ids
+ ):
+ raise ValueError(
+ f"Each list in `bad_words_ids` has to be a list of positive integers, but is {bad_words_ids}."
+ )
+
+ # stores the information about bad words in three tensors:
+ # 1. a rectangular tensor with the forbidden sequences (padded with `-1`), for full data comparisons
+ self.bad_word_seqs_ids = tf.ragged.constant(bad_words_ids).to_tensor(default_value=-1)
+ # 2. a tensor with the unpadded length of each forbidden sequence, for quick length comparisons
+ bad_word_seqs_len = [len(bad_words) for bad_words in bad_words_ids]
+ if any(word_len == 0 for word_len in bad_word_seqs_len):
+ raise ValueError(f"Banned words token sequences {bad_words_ids} cannot have an empty list")
+ self.bad_word_seqs_len = tf.convert_to_tensor(bad_word_seqs_len, dtype=tf.int32)
+ # 3. a tensor containing the last token for each sequence, for easy access to the tokens that may be banned
+ self.seq_forbidden_tokens = tf.convert_to_tensor([bad_words[-1] for bad_words in bad_words_ids])
+
+ def _calc_row_banned_bad_tokens(self, row_input_ids: tf.Tensor) -> tf.Tensor:
+ def _tokens_match(bad_word_seq_number):
+ def _len_one():
+ # If the bad sequence only has one token, always mask it
+ return tf.cond(
+ tf.math.equal(self.bad_word_seqs_len[bad_word_seq_number], 1),
+ lambda: tf.ones((), dtype=tf.bool),
+ _len_greater_than_cur_len,
+ )
+
+ def _len_greater_than_cur_len():
+ # Otherwise, if the bad sequence is longer than the current length they can't ever match
+ return tf.cond(
+ tf.math.greater(self.bad_word_seqs_len[bad_word_seq_number], tf.shape(row_input_ids)[0]),
+ lambda: tf.zeros((), dtype=tf.bool),
+ _match_found,
+ )
+
+ def _match_found():
+ # Finaly, runs the actual comparison. Can only be called if the previous comparisons do not yield
+ # an answer (otherwise we get indexing exceptions)
+ compare_len = self.bad_word_seqs_len[bad_word_seq_number] - 1
+ return tf.cond(
+ tf.math.reduce_all(
+ tf.math.equal(
+ row_input_ids[-compare_len:], self.bad_word_seqs_ids[bad_word_seq_number, :compare_len]
+ )
+ ),
+ lambda: tf.ones((), dtype=tf.bool),
+ lambda: tf.zeros((), dtype=tf.bool),
+ )
+
+ match = _len_one()
+ return match
+
+ # Compares the current row against all bad word sequences, obtaining a mask with the matches.
+ match_mask = tf.map_fn(_tokens_match, tf.range(self.bad_word_seqs_ids.shape[0]), fn_output_signature=tf.bool)
+ row_banned_tokens = self.seq_forbidden_tokens[match_mask]
+ return row_banned_tokens
+
+ def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor:
+ # We want to mask some banned tokens, at a score level. Since the banned tokens depend on the previous
+ # `input_ids`, they may have a different length for each row, and they may even be empty for some rows.
+ # To remain simple and XLA-compatible, we work on a per-row fashion.
+ # TODO (Joao): this function might trigger XLA retracing as `cur_len` increases. Fix it if it becomes
+ # a frequent choke point. (make `cur_len` a tensor?)
+ def _get_row_updated_score(row_inputs: Tuple[tf.Tensor]) -> tf.Tensor:
+ row_input_ids, row_score = row_inputs
+ banned_tokens = self._calc_row_banned_bad_tokens(row_input_ids[:cur_len])
+ banned_tokens_mask = tf.scatter_nd(
+ indices=tf.expand_dims(banned_tokens, axis=-1),
+ updates=tf.ones_like(banned_tokens, dtype=tf.bool),
+ shape=row_score.shape,
+ )
+ row_score = tf.where(banned_tokens_mask, -float("inf"), row_score)
+ return row_score
+
+ scores = tf.map_fn(_get_row_updated_score, (input_ids, scores), fn_output_signature=tf.float32)
+ return scores
+
+
+class TFNoRepeatNGramLogitsProcessor(TFLogitsProcessor):
+ r"""
+ [`TFLogitsProcessor`] that enforces no repetition of n-grams. See
+ [Fairseq](https://github.com/pytorch/fairseq/blob/a07cb6f40480928c9e0548b737aadd36ee66ac76/fairseq/sequence_generator.py#L345).
+
+ Args:
+ ngram_size (`int`):
+ All ngrams of size `ngram_size` can only occur once.
+ """
+
+ def __init__(self, ngram_size: int):
+ if not isinstance(ngram_size, int) or ngram_size <= 0:
+ raise ValueError(f"`ngram_size` has to be a strictly positive integer, but is {ngram_size}")
+ self.ngram_size = ngram_size
+
+ def calc_banned_ngram_tokens(self, input_ids, num_hypos, cur_len):
+ # Copied from fairseq for no_repeat_ngram in beam_search
+ if cur_len + 1 < self.ngram_size:
+ # return no banned tokens if we haven't generated ngram_size tokens yet
+ return [[] for _ in range(num_hypos)]
+ generated_ngrams = [{} for _ in range(num_hypos)]
+ prev_input_ids = input_ids[:, :cur_len]
+ for idx in range(num_hypos):
+ gen_tokens = prev_input_ids[idx].numpy().tolist()
+ generated_ngram = generated_ngrams[idx]
+ for ngram in zip(*[gen_tokens[i:] for i in range(self.ngram_size)]):
+ prev_ngram_tuple = tuple(ngram[:-1])
+ generated_ngram[prev_ngram_tuple] = generated_ngram.get(prev_ngram_tuple, []) + [ngram[-1]]
+
+ def _get_generated_ngrams(hypo_idx):
+ # Before decoding the next token, prevent decoding of ngrams that have already appeared
+ start_idx = cur_len + 1 - self.ngram_size
+ ngram_idx = tuple(prev_input_ids[hypo_idx, start_idx:cur_len].numpy().tolist())
+ return generated_ngrams[hypo_idx].get(ngram_idx, [])
+
+ banned_tokens = [_get_generated_ngrams(hypo_idx) for hypo_idx in range(num_hypos)]
+
+ return banned_tokens
+
+ def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor:
+ # TODO (joao): enable XLA on this logits processor. See discussion and attempts in
+ # https://github.com/huggingface/transformers/pull/16974
+ if not tf.executing_eagerly():
+ raise NotImplementedError("TFNoRepeatNGramLogitsProcessor is only implemented for eager execution.")
+
+ batch_size, vocab_size = scores.shape
+ banned_tokens = self.calc_banned_ngram_tokens(input_ids, batch_size, cur_len)
+
+ # create banned_tokens boolean mask
+ banned_tokens_indices_mask = []
+ for banned_tokens_slice in banned_tokens:
+ banned_tokens_indices_mask.append(
+ [True if token in banned_tokens_slice else False for token in range(vocab_size)]
+ )
+
+ scores = tf.where(tf.convert_to_tensor(banned_tokens_indices_mask, dtype=tf.bool), -float("inf"), scores)
+
+ return scores
+
+
+class TFForcedBOSTokenLogitsProcessor(TFLogitsProcessor):
+ r"""
+ [`TFLogitsProcessor`] that enforces the specified token as the first generated token.
+
+ Args:
+ bos_token_id (`int`):
+ The id of the token to force as the first generated token.
+ """
+
+ def __init__(self, bos_token_id: int):
+ if bos_token_id < 0:
+ raise ValueError(f"The forced bos token id must be a non-negative integer, got {bos_token_id}")
+ self.bos_token_id = bos_token_id
+
+ def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor:
+ if cur_len == 1:
+ batch_size, num_tokens = scores.shape
+ # sets the score to 0 in the bos_token_id column
+ scores = tf.zeros((batch_size, 1))
+ # sets the score to -inf everywhere else
+ if self.bos_token_id > 0:
+ scores = tf.concat((tf.broadcast_to(-float("inf"), (batch_size, self.bos_token_id)), scores), axis=-1)
+ if self.bos_token_id < (num_tokens - 1):
+ scores = tf.concat(
+ (scores, tf.broadcast_to(-float("inf"), (batch_size, (num_tokens - 1) - self.bos_token_id))),
+ axis=-1,
+ )
+ return scores
+
+
+class TFForcedEOSTokenLogitsProcessor(TFLogitsProcessor):
+ r"""
+ [`TFLogitsProcessor`] that enforces the specified token as the last generated token when `max_length` is reached.
+
+ Args:
+ max_length (`int`):
+ The maximum length of the sequence to be generated.
+ eos_token_id (`int`):
+ The id of the token to force as the last generated token when `max_length` is reached.
+ """
+
+ def __init__(self, max_length: int, eos_token_id: int):
+ self.max_length = max_length
+ if eos_token_id < 0:
+ raise ValueError(f"The forced eos token id must be a non-negative integer, got {eos_token_id}")
+ self.eos_token_id = eos_token_id
+
+ def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor:
+ if cur_len == self.max_length - 1:
+ batch_size, num_tokens = scores.shape
+ # sets the score to 0 in the eos_token_id column
+ scores = tf.zeros((batch_size, 1))
+ # sets the score to -inf everywhere else
+ if self.eos_token_id > 0:
+ scores = tf.concat((tf.broadcast_to(-float("inf"), (batch_size, self.eos_token_id)), scores), axis=-1)
+ if self.eos_token_id < (num_tokens - 1):
+ scores = tf.concat(
+ (scores, tf.broadcast_to(-float("inf"), (batch_size, (num_tokens - 1) - self.eos_token_id))),
+ axis=-1,
+ )
+ return scores
+
+
+class TFSuppressTokensAtBeginLogitsProcessor(TFLogitsProcessor):
+ r"""
+ [`TFSuppressTokensAtBeginLogitsProcessor`] suppresses a list of tokens as soon as the `generate` function starts
+ generating using `begin_index` tokens. This should ensure that the tokens defined by `begin_suppress_tokens` at not
+ sampled at the beginning of the generation.
+ """
+
+ def __init__(self, begin_suppress_tokens, begin_index):
+ self.begin_suppress_tokens = list(begin_suppress_tokens)
+ self.begin_index = begin_index
+
+ def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor:
+ suppressed_indices = []
+ for token in self.begin_suppress_tokens:
+ if token < scores.shape[-1]: # to ensure we don't go beyond the vocab size
+ suppressed_indices.extend([[i, token] for i in range(scores.shape[0])])
+
+ if len(suppressed_indices) > 0:
+ scores = tf.cond(
+ tf.equal(cur_len, self.begin_index),
+ lambda: tf.tensor_scatter_nd_update(
+ scores,
+ indices=suppressed_indices,
+ updates=[-float("inf") for _ in range(scores.shape[0] * len(self.begin_suppress_tokens))],
+ ),
+ lambda: scores,
+ )
+ return scores
+
+
+class TFSuppressTokensLogitsProcessor(TFLogitsProcessor):
+ r"""This processor can be used to suppress a list of tokens. The processor will set their log probs to `-inf` so that they
+ are not sampled."""
+
+ def __init__(self, suppress_tokens):
+ self.suppress_tokens = list(suppress_tokens)
+
+ def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor:
+ suppressed_indices = []
+ for token in self.suppress_tokens:
+ if token < scores.shape[-1]: # to ensure we don't go beyond the vocab size
+ suppressed_indices.extend([[i, token] for i in range(scores.shape[0])])
+
+ if len(suppressed_indices) > 0:
+ scores = tf.tensor_scatter_nd_update(
+ scores,
+ indices=[[i, token] for i in range(scores.shape[0]) for token in self.suppress_tokens],
+ updates=[-float("inf") for _ in range(scores.shape[0] * len(self.suppress_tokens))],
+ )
+ return scores
+
+
+class TFForceTokensLogitsProcessor(TFLogitsProcessor):
+ r"""This processor takes a list of pairs of integers which indicates a mapping from generation indices to token
+ indices that will be forced before sampling. The processor will set their log probs to `0` and all other tokens to
+ `-inf` so that they are sampled at their corresponding index."""
+
+ def __init__(self, force_token_map: List[List[int]]):
+ force_token_map = dict(force_token_map)
+ # Converts the dictionary of format {index: token} containing the tokens to be forced to an array, where the
+ # index of the array corresponds to the index of the token to be forced, for XLA compatibility.
+ # Indexes without forced tokens will have an negative value.
+ force_token_array = np.ones((max(force_token_map.keys()) + 1), dtype=np.int32) * -1
+ for index, token in force_token_map.items():
+ if token is not None:
+ force_token_array[index] = token
+ self.force_token_array = tf.convert_to_tensor(force_token_array, dtype=tf.int32)
+
+ def __call__(self, input_ids: tf.Tensor, scores: tf.Tensor, cur_len: int) -> tf.Tensor:
+ def _force_token(generation_idx):
+ batch_size = scores.shape[0]
+ current_token = self.force_token_array[generation_idx]
+
+ new_scores = tf.zeros_like(scores, dtype=scores.dtype) + tf.constant([scores.dtype.min])
+ indices = tf.stack((tf.range(batch_size), tf.tile([current_token], [batch_size])), axis=1)
+ updates = tf.zeros((batch_size,), dtype=scores.dtype)
+ new_scores = tf.tensor_scatter_nd_update(new_scores, indices, updates)
+ return new_scores
+
+ scores = tf.cond(
+ tf.greater_equal(cur_len, tf.shape(self.force_token_array)[0]),
+ # If the current length is geq than the length of force_token_array, the processor does nothing.
+ lambda: tf.identity(scores),
+ # Otherwise, it may force a certain token.
+ lambda: tf.cond(
+ tf.greater_equal(self.force_token_array[cur_len], 0),
+ # Only valid (positive) tokens are forced
+ lambda: _force_token(cur_len),
+ # Otherwise, the processor does nothing.
+ lambda: scores,
+ ),
+ )
+ return scores
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/tf_utils.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/tf_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..c3770cb1237f3c6ea3517bda3dd29b2b6f44f7de
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/tf_utils.py
@@ -0,0 +1,294 @@
+# Copyright 2022 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.
+
+from typing import Optional, Union
+
+import numpy as np
+import tensorflow as tf
+
+from .feature_extraction_utils import BatchFeature
+from .tokenization_utils_base import BatchEncoding
+from .utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+def shape_list(tensor: Union[tf.Tensor, np.ndarray]) -> list[int]:
+ """
+ Deal with dynamic shape in tensorflow cleanly.
+
+ Args:
+ tensor (`tf.Tensor` or `np.ndarray`): The tensor we want the shape of.
+
+ Returns:
+ `List[int]`: The shape of the tensor as a list.
+ """
+ if isinstance(tensor, np.ndarray):
+ return list(tensor.shape)
+
+ dynamic = tf.shape(tensor)
+
+ if tensor.shape == tf.TensorShape(None):
+ return dynamic
+
+ static = tensor.shape.as_list()
+
+ return [dynamic[i] if s is None else s for i, s in enumerate(static)]
+
+
+def stable_softmax(logits: tf.Tensor, axis: Optional[int] = None, name: Optional[str] = None) -> tf.Tensor:
+ """
+ Stable wrapper that returns the same output as `tf.nn.softmax`, but that works reliably with XLA on CPU. It is
+ meant as a workaround for the [following issue](https://github.com/tensorflow/tensorflow/issues/55682), and will be
+ removed after it gets fixed. The arguments and outputs are the same as `tf.nn.softmax`, and relies on the fact that
+ `softmax(x) = softmax(x + c)` (see https://ogunlao.github.io/2020/04/26/you_dont_really_know_softmax.html).
+
+ Args:
+ logits (`tf.Tensor`):
+ Must be one of the following types: half, float32, float64.
+ axis (`int`, *optional*):
+ The dimension softmax would be performed on. The default is -1 which indicates the last dimension.
+ name (`str`, *optional*):
+ A name for the operation.
+
+ Returns:
+ `tf.Tensor`:
+ A Tensor. Has the same type and shape as logits.
+ """
+ # TODO: When the issue linked above gets sorted, add a check on TF version here and use the original function if
+ # it has the fix. After we drop the support for unfixed versions, remove this function.
+ return tf.nn.softmax(logits=logits + 1e-9, axis=axis, name=name)
+
+
+def functional_layernorm(inputs, weight, bias, epsilon=1e-5, axis=-1):
+ # This is a very simplified functional layernorm, designed to duplicate
+ # the functionality of PyTorch nn.functional.layer_norm when this is needed to port
+ # models in Transformers.
+
+ if weight.shape.rank != 1 or bias.shape.rank != 1 or not isinstance(axis, int):
+ raise NotImplementedError("Only 1D weight and bias tensors are supported for now, with only a single axis.")
+
+ # Get mean and variance on the axis to be normalized
+ mean, variance = tf.nn.moments(inputs, axes=[axis], keepdims=True)
+
+ if axis != -1:
+ # Reshape scale and weight to have the same rank as inputs, but with 1 dimensions
+ # on every dimension except axis
+ shape = [1] * inputs.shape.rank
+ shape[axis] = shape_list(inputs)[axis]
+ weight = tf.reshape(weight, shape)
+ bias = tf.reshape(bias, shape)
+
+ # Compute layer normalization using the batch_normalization
+ # function.
+ outputs = tf.nn.batch_normalization(
+ inputs,
+ mean,
+ variance,
+ offset=bias,
+ scale=weight,
+ variance_epsilon=epsilon,
+ )
+ return outputs
+
+
+def scaled_dot_product_attention(
+ query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale: Optional[float] = None
+):
+ """TF equivalent for torch's nn.functional.scaled_dot_product_attention"""
+ if dropout_p != 0.0:
+ raise ValueError(
+ "Dropout is not supported in this implementation - file an issue "
+ "with Transformers and ping @Rocketknight1 if you need it for a port!"
+ )
+ if is_causal and attn_mask is not None:
+ raise ValueError("You cannot specify an attn_mask and is_causal at the same time!")
+ if is_causal:
+ attn_mask = tf.ones((tf.shape(query)[-2], tf.shape(key)[-2]), dtype=tf.int32)
+ attn_mask = tf.experimental.numpy.tril(attn_mask, k=0)
+ if attn_mask is not None and (attn_mask.dtype.is_integer or attn_mask.dtype.is_bool):
+ # Convert boolean mask to a negative logit bias
+ attn_mask = tf.where(attn_mask > 0, tf.cast(0.0, query.dtype), tf.cast(-1000.0, query.dtype))
+ logits = tf.einsum("...qd, ...kd -> ...qk", query, key)
+ if scale is None:
+ scale = tf.cast(tf.shape(key)[-1], logits.dtype) ** -0.5
+ logits *= scale # scale by 1/sqrt(key_dim)
+ if attn_mask is not None:
+ logits += attn_mask
+ probs = tf.nn.softmax(logits)
+ return probs @ value
+
+
+def flatten(input, start_dim=0, end_dim=-1):
+ # Replicates the behavior of torch.flatten in TF
+
+ # If end_dim or start_dim is negative, count them from the end
+ if end_dim < 0:
+ end_dim += input.shape.rank
+ if start_dim < 0:
+ start_dim += input.shape.rank
+
+ if start_dim == end_dim:
+ return input
+
+ in_shape = tf.shape(input)
+ flattened_dim = tf.math.reduce_prod(in_shape[start_dim : end_dim + 1])
+ out_shape = tf.concat([in_shape[:start_dim], [flattened_dim], in_shape[end_dim + 1 :]], axis=0)
+ return tf.reshape(input, out_shape)
+
+
+def invert_attention_mask(encoder_attention_mask: tf.Tensor) -> tf.Tensor:
+ """
+ Invert an attention mask (e.g., switches 0. and 1.).
+
+ Args:
+ encoder_attention_mask (`torch.Tensor`): An attention mask.
+
+ Returns:
+ `tf.Tensor`: The inverted attention mask.
+ """
+ if not isinstance(encoder_attention_mask, tf.Tensor):
+ encoder_attention_mask = tf.convert_to_tensor(encoder_attention_mask) # Catches stray NumPy inputs
+ if encoder_attention_mask.shape.rank == 3:
+ encoder_extended_attention_mask = encoder_attention_mask[:, None, :, :]
+ if encoder_attention_mask.shape.rank == 2:
+ encoder_extended_attention_mask = encoder_attention_mask[:, None, None, :]
+ # T5 has a mask that can compare sequence ids, we can simulate this here with this transposition
+ # Cf. https://github.com/tensorflow/mesh/blob/8d2465e9bc93129b913b5ccc6a59aa97abd96ec6/mesh_tensorflow
+ # /transformer/transformer_layers.py#L270
+ # encoder_extended_attention_mask = (encoder_extended_attention_mask ==
+ # encoder_extended_attention_mask.transpose(-1, -2))
+ encoder_extended_attention_mask = (
+ tf.cast(1, encoder_attention_mask.dtype) - encoder_extended_attention_mask
+ ) * encoder_extended_attention_mask.dtype.min
+
+ return encoder_extended_attention_mask
+
+
+def check_embeddings_within_bounds(tensor: tf.Tensor, embed_dim: int, tensor_name: str = "input_ids") -> None:
+ """
+ `tf.gather`, on which TF embedding layers are based, won't check positive out of bound indices on GPU, returning
+ zeros instead. This function adds a check against that dangerous silent behavior.
+
+ Args:
+ tensor (`tf.Tensor`): The tensor of indices to check.
+ embed_dim (`int`): The embedding dimension.
+ tensor_name (`str`, *optional*): The name of the tensor to use in the error message.
+ """
+ tf.debugging.assert_less(
+ tensor,
+ tf.cast(embed_dim, dtype=tensor.dtype),
+ message=(
+ f"The maximum value of {tensor_name} ({tf.math.reduce_max(tensor)}) must be smaller than the embedding "
+ f"layer's input dimension ({embed_dim}). The likely cause is some problem at tokenization time."
+ ),
+ )
+
+
+def save_attributes_to_hdf5_group(group, name, data):
+ """Saves attributes (data) of the specified name into the HDF5 group.
+
+ This method deals with an inherent problem of HDF5 file which is not able to store data larger than
+ HDF5_OBJECT_HEADER_LIMIT bytes.
+
+ Args:
+ group: A pointer to a HDF5 group.
+ name: A name of the attributes to save.
+ data: Attributes data to store.
+
+ Raises:
+ RuntimeError: If any single attribute is too large to be saved.
+
+ Copied from Keras to Transformers to avoid versioning issues.
+ """
+ HDF5_OBJECT_HEADER_LIMIT = 64512
+ # Check that no item in `data` is larger than `HDF5_OBJECT_HEADER_LIMIT`
+ # because in that case even chunking the array would not make the saving
+ # possible.
+ bad_attributes = [x for x in data if len(x) > HDF5_OBJECT_HEADER_LIMIT]
+
+ # Expecting this to never be true.
+ if bad_attributes:
+ raise RuntimeError(
+ "The following attributes cannot be saved to HDF5 file because "
+ f"they are larger than {HDF5_OBJECT_HEADER_LIMIT} "
+ f"bytes: {bad_attributes}"
+ )
+
+ data_npy = np.asarray(data)
+
+ num_chunks = 1
+ chunked_data = np.array_split(data_npy, num_chunks)
+
+ # This will never loop forever thanks to the test above.
+ while any(x.nbytes > HDF5_OBJECT_HEADER_LIMIT for x in chunked_data):
+ num_chunks += 1
+ chunked_data = np.array_split(data_npy, num_chunks)
+
+ if num_chunks > 1:
+ for chunk_id, chunk_data in enumerate(chunked_data):
+ group.attrs["%s%d" % (name, chunk_id)] = chunk_data
+ else:
+ group.attrs[name] = data
+
+
+def load_attributes_from_hdf5_group(group, name):
+ """Loads attributes of the specified name from the HDF5 group.
+
+ This method deals with an inherent problem of HDF5 file which is not able to store data larger than
+ HDF5_OBJECT_HEADER_LIMIT bytes.
+
+ Args:
+ group: A pointer to a HDF5 group.
+ name: A name of the attributes to load.
+
+ Returns:
+ data: Attributes data.
+
+ Copied from Keras to Transformers to avoid versioning issues.
+ """
+ if name in group.attrs:
+ data = [n.decode("utf8") if hasattr(n, "decode") else n for n in group.attrs[name]]
+ else:
+ data = []
+ chunk_id = 0
+ while "%s%d" % (name, chunk_id) in group.attrs:
+ data.extend(
+ [n.decode("utf8") if hasattr(n, "decode") else n for n in group.attrs["%s%d" % (name, chunk_id)]]
+ )
+ chunk_id += 1
+ return data
+
+
+def expand_1d(data):
+ """Expands 1-dimensional `Tensor`s into 2-dimensional `Tensor`s.
+ Copied from Keras to here to avoid versioning issues."""
+
+ def _expand_single_1d_tensor(t):
+ if isinstance(t, tf.Tensor) and t.shape.rank == 1:
+ return tf.expand_dims(t, axis=-1)
+ return t
+
+ return tf.nest.map_structure(_expand_single_1d_tensor, data)
+
+
+def convert_batch_encoding(*args, **kwargs):
+ # Convert HF BatchEncoding/BatchFeature objects in the inputs to dicts that Keras understands
+ if args and isinstance(args[0], (BatchEncoding, BatchFeature)):
+ args = list(args)
+ args[0] = dict(args[0])
+ elif "x" in kwargs and isinstance(kwargs["x"], (BatchEncoding, BatchFeature)):
+ kwargs["x"] = dict(kwargs["x"])
+ return args, kwargs
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/tokenization_auto.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/tokenization_auto.py
new file mode 100644
index 0000000000000000000000000000000000000000..eda7356a60834881116acc2ed02eff39f537dc39
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/tokenization_auto.py
@@ -0,0 +1,1085 @@
+# coding=utf-8
+# Copyright 2018 The HuggingFace Inc. 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.
+"""Auto Tokenizer class."""
+
+import importlib
+import json
+import os
+import warnings
+from collections import OrderedDict
+from typing import TYPE_CHECKING, Dict, Optional, Tuple, Union
+
+from ...configuration_utils import PretrainedConfig
+from ...dynamic_module_utils import get_class_from_dynamic_module, resolve_trust_remote_code
+from ...modeling_gguf_pytorch_utils import load_gguf_checkpoint
+from ...tokenization_utils import PreTrainedTokenizer
+from ...tokenization_utils_base import TOKENIZER_CONFIG_FILE
+from ...utils import (
+ cached_file,
+ extract_commit_hash,
+ is_g2p_en_available,
+ is_sentencepiece_available,
+ is_tokenizers_available,
+ logging,
+)
+from ..encoder_decoder import EncoderDecoderConfig
+from .auto_factory import _LazyAutoMapping
+from .configuration_auto import (
+ CONFIG_MAPPING_NAMES,
+ AutoConfig,
+ config_class_to_model_type,
+ model_type_to_module_name,
+ replace_list_option_in_docstrings,
+)
+
+
+if is_tokenizers_available():
+ from ...tokenization_utils_fast import PreTrainedTokenizerFast
+else:
+ PreTrainedTokenizerFast = None
+
+
+logger = logging.get_logger(__name__)
+
+if TYPE_CHECKING:
+ # This significantly improves completion suggestion performance when
+ # the transformers package is used with Microsoft's Pylance language server.
+ TOKENIZER_MAPPING_NAMES: OrderedDict[str, Tuple[Optional[str], Optional[str]]] = OrderedDict()
+else:
+ TOKENIZER_MAPPING_NAMES = OrderedDict(
+ [
+ (
+ "albert",
+ (
+ "AlbertTokenizer" if is_sentencepiece_available() else None,
+ "AlbertTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ ("align", ("BertTokenizer", "BertTokenizerFast" if is_tokenizers_available() else None)),
+ ("aria", ("LlamaTokenizer", "LlamaTokenizerFast" if is_tokenizers_available() else None)),
+ ("aya_vision", (None, "CohereTokenizerFast" if is_tokenizers_available() else None)),
+ ("bark", ("BertTokenizer", "BertTokenizerFast" if is_tokenizers_available() else None)),
+ ("bart", ("BartTokenizer", "BartTokenizerFast")),
+ (
+ "barthez",
+ (
+ "BarthezTokenizer" if is_sentencepiece_available() else None,
+ "BarthezTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ ("bartpho", ("BartphoTokenizer", None)),
+ ("bert", ("BertTokenizer", "BertTokenizerFast" if is_tokenizers_available() else None)),
+ ("bert-generation", ("BertGenerationTokenizer" if is_sentencepiece_available() else None, None)),
+ ("bert-japanese", ("BertJapaneseTokenizer", None)),
+ ("bertweet", ("BertweetTokenizer", None)),
+ (
+ "big_bird",
+ (
+ "BigBirdTokenizer" if is_sentencepiece_available() else None,
+ "BigBirdTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ ("bigbird_pegasus", ("PegasusTokenizer", "PegasusTokenizerFast" if is_tokenizers_available() else None)),
+ ("biogpt", ("BioGptTokenizer", None)),
+ ("blenderbot", ("BlenderbotTokenizer", "BlenderbotTokenizerFast")),
+ ("blenderbot-small", ("BlenderbotSmallTokenizer", None)),
+ ("blip", ("BertTokenizer", "BertTokenizerFast" if is_tokenizers_available() else None)),
+ ("blip-2", ("GPT2Tokenizer", "GPT2TokenizerFast" if is_tokenizers_available() else None)),
+ ("bloom", (None, "BloomTokenizerFast" if is_tokenizers_available() else None)),
+ ("bridgetower", ("RobertaTokenizer", "RobertaTokenizerFast" if is_tokenizers_available() else None)),
+ ("bros", ("BertTokenizer", "BertTokenizerFast" if is_tokenizers_available() else None)),
+ ("byt5", ("ByT5Tokenizer", None)),
+ (
+ "camembert",
+ (
+ "CamembertTokenizer" if is_sentencepiece_available() else None,
+ "CamembertTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ ("canine", ("CanineTokenizer", None)),
+ (
+ "chameleon",
+ (
+ "LlamaTokenizer" if is_sentencepiece_available() else None,
+ "LlamaTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ ("chinese_clip", ("BertTokenizer", "BertTokenizerFast" if is_tokenizers_available() else None)),
+ (
+ "clap",
+ (
+ "RobertaTokenizer",
+ "RobertaTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ (
+ "clip",
+ (
+ "CLIPTokenizer",
+ "CLIPTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ (
+ "clipseg",
+ (
+ "CLIPTokenizer",
+ "CLIPTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ ("clvp", ("ClvpTokenizer", None)),
+ (
+ "code_llama",
+ (
+ "CodeLlamaTokenizer" if is_sentencepiece_available() else None,
+ "CodeLlamaTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ ("codegen", ("CodeGenTokenizer", "CodeGenTokenizerFast" if is_tokenizers_available() else None)),
+ ("cohere", (None, "CohereTokenizerFast" if is_tokenizers_available() else None)),
+ ("cohere2", (None, "CohereTokenizerFast" if is_tokenizers_available() else None)),
+ ("colpali", ("LlamaTokenizer", "LlamaTokenizerFast" if is_tokenizers_available() else None)),
+ ("convbert", ("ConvBertTokenizer", "ConvBertTokenizerFast" if is_tokenizers_available() else None)),
+ (
+ "cpm",
+ (
+ "CpmTokenizer" if is_sentencepiece_available() else None,
+ "CpmTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ ("cpmant", ("CpmAntTokenizer", None)),
+ ("ctrl", ("CTRLTokenizer", None)),
+ ("data2vec-audio", ("Wav2Vec2CTCTokenizer", None)),
+ ("data2vec-text", ("RobertaTokenizer", "RobertaTokenizerFast" if is_tokenizers_available() else None)),
+ ("dbrx", ("GPT2Tokenizer", "GPT2TokenizerFast" if is_tokenizers_available() else None)),
+ ("deberta", ("DebertaTokenizer", "DebertaTokenizerFast" if is_tokenizers_available() else None)),
+ (
+ "deberta-v2",
+ (
+ "DebertaV2Tokenizer" if is_sentencepiece_available() else None,
+ "DebertaV2TokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ (
+ "deepseek_v3",
+ (
+ "LlamaTokenizer" if is_sentencepiece_available() else None,
+ "LlamaTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ (
+ "diffllama",
+ (
+ "LlamaTokenizer" if is_sentencepiece_available() else None,
+ "LlamaTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ ("distilbert", ("DistilBertTokenizer", "DistilBertTokenizerFast" if is_tokenizers_available() else None)),
+ (
+ "dpr",
+ (
+ "DPRQuestionEncoderTokenizer",
+ "DPRQuestionEncoderTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ ("electra", ("ElectraTokenizer", "ElectraTokenizerFast" if is_tokenizers_available() else None)),
+ ("emu3", ("GPT2Tokenizer", "GPT2TokenizerFast" if is_tokenizers_available() else None)),
+ ("ernie", ("BertTokenizer", "BertTokenizerFast" if is_tokenizers_available() else None)),
+ ("ernie_m", ("ErnieMTokenizer" if is_sentencepiece_available() else None, None)),
+ ("esm", ("EsmTokenizer", None)),
+ ("falcon", (None, "PreTrainedTokenizerFast" if is_tokenizers_available() else None)),
+ ("falcon_mamba", (None, "GPTNeoXTokenizerFast" if is_tokenizers_available() else None)),
+ (
+ "fastspeech2_conformer",
+ ("FastSpeech2ConformerTokenizer" if is_g2p_en_available() else None, None),
+ ),
+ ("flaubert", ("FlaubertTokenizer", None)),
+ ("fnet", ("FNetTokenizer", "FNetTokenizerFast" if is_tokenizers_available() else None)),
+ ("fsmt", ("FSMTTokenizer", None)),
+ ("funnel", ("FunnelTokenizer", "FunnelTokenizerFast" if is_tokenizers_available() else None)),
+ (
+ "gemma",
+ (
+ "GemmaTokenizer" if is_sentencepiece_available() else None,
+ "GemmaTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ (
+ "gemma2",
+ (
+ "GemmaTokenizer" if is_sentencepiece_available() else None,
+ "GemmaTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ (
+ "gemma3",
+ (
+ "GemmaTokenizer" if is_sentencepiece_available() else None,
+ "GemmaTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ (
+ "gemma3_text",
+ (
+ "GemmaTokenizer" if is_sentencepiece_available() else None,
+ "GemmaTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ ("git", ("BertTokenizer", "BertTokenizerFast" if is_tokenizers_available() else None)),
+ ("glm", (None, "PreTrainedTokenizerFast" if is_tokenizers_available() else None)),
+ ("glm4", (None, "PreTrainedTokenizerFast" if is_tokenizers_available() else None)),
+ ("gpt-sw3", ("GPTSw3Tokenizer" if is_sentencepiece_available() else None, None)),
+ ("gpt2", ("GPT2Tokenizer", "GPT2TokenizerFast" if is_tokenizers_available() else None)),
+ ("gpt_bigcode", ("GPT2Tokenizer", "GPT2TokenizerFast" if is_tokenizers_available() else None)),
+ ("gpt_neo", ("GPT2Tokenizer", "GPT2TokenizerFast" if is_tokenizers_available() else None)),
+ ("gpt_neox", (None, "GPTNeoXTokenizerFast" if is_tokenizers_available() else None)),
+ ("gpt_neox_japanese", ("GPTNeoXJapaneseTokenizer", None)),
+ ("gptj", ("GPT2Tokenizer", "GPT2TokenizerFast" if is_tokenizers_available() else None)),
+ ("gptsan-japanese", ("GPTSanJapaneseTokenizer", None)),
+ ("grounding-dino", ("BertTokenizer", "BertTokenizerFast" if is_tokenizers_available() else None)),
+ ("groupvit", ("CLIPTokenizer", "CLIPTokenizerFast" if is_tokenizers_available() else None)),
+ ("helium", (None, "PreTrainedTokenizerFast" if is_tokenizers_available() else None)),
+ ("herbert", ("HerbertTokenizer", "HerbertTokenizerFast" if is_tokenizers_available() else None)),
+ ("hubert", ("Wav2Vec2CTCTokenizer", None)),
+ ("ibert", ("RobertaTokenizer", "RobertaTokenizerFast" if is_tokenizers_available() else None)),
+ ("idefics", (None, "LlamaTokenizerFast" if is_tokenizers_available() else None)),
+ ("idefics2", ("LlamaTokenizer", "LlamaTokenizerFast" if is_tokenizers_available() else None)),
+ ("idefics3", ("LlamaTokenizer", "LlamaTokenizerFast" if is_tokenizers_available() else None)),
+ ("instructblip", ("GPT2Tokenizer", "GPT2TokenizerFast" if is_tokenizers_available() else None)),
+ ("instructblipvideo", ("GPT2Tokenizer", "GPT2TokenizerFast" if is_tokenizers_available() else None)),
+ (
+ "jamba",
+ (
+ "LlamaTokenizer" if is_sentencepiece_available() else None,
+ "LlamaTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ (
+ "jetmoe",
+ (
+ "LlamaTokenizer" if is_sentencepiece_available() else None,
+ "LlamaTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ ("jukebox", ("JukeboxTokenizer", None)),
+ (
+ "kosmos-2",
+ (
+ "XLMRobertaTokenizer" if is_sentencepiece_available() else None,
+ "XLMRobertaTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ ("layoutlm", ("LayoutLMTokenizer", "LayoutLMTokenizerFast" if is_tokenizers_available() else None)),
+ ("layoutlmv2", ("LayoutLMv2Tokenizer", "LayoutLMv2TokenizerFast" if is_tokenizers_available() else None)),
+ ("layoutlmv3", ("LayoutLMv3Tokenizer", "LayoutLMv3TokenizerFast" if is_tokenizers_available() else None)),
+ ("layoutxlm", ("LayoutXLMTokenizer", "LayoutXLMTokenizerFast" if is_tokenizers_available() else None)),
+ ("led", ("LEDTokenizer", "LEDTokenizerFast" if is_tokenizers_available() else None)),
+ ("lilt", ("LayoutLMv3Tokenizer", "LayoutLMv3TokenizerFast" if is_tokenizers_available() else None)),
+ (
+ "llama",
+ (
+ "LlamaTokenizer" if is_sentencepiece_available() else None,
+ "LlamaTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ (
+ "llama4",
+ (
+ "LlamaTokenizer" if is_sentencepiece_available() else None,
+ "LlamaTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ (
+ "llama4_text",
+ (
+ "LlamaTokenizer" if is_sentencepiece_available() else None,
+ "LlamaTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ ("llava", ("LlamaTokenizer", "LlamaTokenizerFast" if is_tokenizers_available() else None)),
+ ("llava_next", ("LlamaTokenizer", "LlamaTokenizerFast" if is_tokenizers_available() else None)),
+ ("llava_next_video", ("LlamaTokenizer", "LlamaTokenizerFast" if is_tokenizers_available() else None)),
+ ("llava_onevision", ("LlamaTokenizer", "LlamaTokenizerFast" if is_tokenizers_available() else None)),
+ ("longformer", ("LongformerTokenizer", "LongformerTokenizerFast" if is_tokenizers_available() else None)),
+ (
+ "longt5",
+ (
+ "T5Tokenizer" if is_sentencepiece_available() else None,
+ "T5TokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ ("luke", ("LukeTokenizer", None)),
+ ("lxmert", ("LxmertTokenizer", "LxmertTokenizerFast" if is_tokenizers_available() else None)),
+ ("m2m_100", ("M2M100Tokenizer" if is_sentencepiece_available() else None, None)),
+ ("mamba", (None, "GPTNeoXTokenizerFast" if is_tokenizers_available() else None)),
+ ("mamba2", (None, "GPTNeoXTokenizerFast" if is_tokenizers_available() else None)),
+ ("marian", ("MarianTokenizer" if is_sentencepiece_available() else None, None)),
+ (
+ "mbart",
+ (
+ "MBartTokenizer" if is_sentencepiece_available() else None,
+ "MBartTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ (
+ "mbart50",
+ (
+ "MBart50Tokenizer" if is_sentencepiece_available() else None,
+ "MBart50TokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ ("mega", ("RobertaTokenizer", "RobertaTokenizerFast" if is_tokenizers_available() else None)),
+ ("megatron-bert", ("BertTokenizer", "BertTokenizerFast" if is_tokenizers_available() else None)),
+ ("mgp-str", ("MgpstrTokenizer", None)),
+ (
+ "mistral",
+ (
+ "LlamaTokenizer" if is_sentencepiece_available() else None,
+ "LlamaTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ (
+ "mixtral",
+ (
+ "LlamaTokenizer" if is_sentencepiece_available() else None,
+ "LlamaTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ ("mllama", ("LlamaTokenizer", "LlamaTokenizerFast" if is_tokenizers_available() else None)),
+ ("mluke", ("MLukeTokenizer" if is_sentencepiece_available() else None, None)),
+ ("mobilebert", ("MobileBertTokenizer", "MobileBertTokenizerFast" if is_tokenizers_available() else None)),
+ ("modernbert", (None, "PreTrainedTokenizerFast" if is_tokenizers_available() else None)),
+ ("moonshine", (None, "PreTrainedTokenizerFast" if is_tokenizers_available() else None)),
+ ("moshi", (None, "PreTrainedTokenizerFast" if is_tokenizers_available() else None)),
+ ("mpnet", ("MPNetTokenizer", "MPNetTokenizerFast" if is_tokenizers_available() else None)),
+ ("mpt", (None, "GPTNeoXTokenizerFast" if is_tokenizers_available() else None)),
+ ("mra", ("RobertaTokenizer", "RobertaTokenizerFast" if is_tokenizers_available() else None)),
+ (
+ "mt5",
+ (
+ "MT5Tokenizer" if is_sentencepiece_available() else None,
+ "MT5TokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ ("musicgen", ("T5Tokenizer", "T5TokenizerFast" if is_tokenizers_available() else None)),
+ ("musicgen_melody", ("T5Tokenizer", "T5TokenizerFast" if is_tokenizers_available() else None)),
+ ("mvp", ("MvpTokenizer", "MvpTokenizerFast" if is_tokenizers_available() else None)),
+ ("myt5", ("MyT5Tokenizer", None)),
+ ("nemotron", (None, "PreTrainedTokenizerFast" if is_tokenizers_available() else None)),
+ ("nezha", ("BertTokenizer", "BertTokenizerFast" if is_tokenizers_available() else None)),
+ (
+ "nllb",
+ (
+ "NllbTokenizer" if is_sentencepiece_available() else None,
+ "NllbTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ (
+ "nllb-moe",
+ (
+ "NllbTokenizer" if is_sentencepiece_available() else None,
+ "NllbTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ (
+ "nystromformer",
+ (
+ "AlbertTokenizer" if is_sentencepiece_available() else None,
+ "AlbertTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ ("olmo", (None, "GPTNeoXTokenizerFast" if is_tokenizers_available() else None)),
+ ("olmo2", (None, "GPTNeoXTokenizerFast" if is_tokenizers_available() else None)),
+ ("olmoe", (None, "GPTNeoXTokenizerFast" if is_tokenizers_available() else None)),
+ (
+ "omdet-turbo",
+ ("CLIPTokenizer", "CLIPTokenizerFast" if is_tokenizers_available() else None),
+ ),
+ ("oneformer", ("CLIPTokenizer", "CLIPTokenizerFast" if is_tokenizers_available() else None)),
+ (
+ "openai-gpt",
+ ("OpenAIGPTTokenizer", "OpenAIGPTTokenizerFast" if is_tokenizers_available() else None),
+ ),
+ ("opt", ("GPT2Tokenizer", "GPT2TokenizerFast" if is_tokenizers_available() else None)),
+ ("owlv2", ("CLIPTokenizer", "CLIPTokenizerFast" if is_tokenizers_available() else None)),
+ ("owlvit", ("CLIPTokenizer", "CLIPTokenizerFast" if is_tokenizers_available() else None)),
+ ("paligemma", ("LlamaTokenizer", "LlamaTokenizerFast" if is_tokenizers_available() else None)),
+ (
+ "pegasus",
+ (
+ "PegasusTokenizer" if is_sentencepiece_available() else None,
+ "PegasusTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ (
+ "pegasus_x",
+ (
+ "PegasusTokenizer" if is_sentencepiece_available() else None,
+ "PegasusTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ (
+ "perceiver",
+ (
+ "PerceiverTokenizer",
+ None,
+ ),
+ ),
+ (
+ "persimmon",
+ (
+ "LlamaTokenizer" if is_sentencepiece_available() else None,
+ "LlamaTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ ("phi", ("CodeGenTokenizer", "CodeGenTokenizerFast" if is_tokenizers_available() else None)),
+ ("phi3", ("LlamaTokenizer", "LlamaTokenizerFast" if is_tokenizers_available() else None)),
+ ("phimoe", ("LlamaTokenizer", "LlamaTokenizerFast" if is_tokenizers_available() else None)),
+ ("phobert", ("PhobertTokenizer", None)),
+ ("pix2struct", ("T5Tokenizer", "T5TokenizerFast" if is_tokenizers_available() else None)),
+ ("pixtral", (None, "PreTrainedTokenizerFast" if is_tokenizers_available() else None)),
+ ("plbart", ("PLBartTokenizer" if is_sentencepiece_available() else None, None)),
+ ("prophetnet", ("ProphetNetTokenizer", None)),
+ ("qdqbert", ("BertTokenizer", "BertTokenizerFast" if is_tokenizers_available() else None)),
+ (
+ "qwen2",
+ (
+ "Qwen2Tokenizer",
+ "Qwen2TokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ ("qwen2_5_vl", ("Qwen2Tokenizer", "Qwen2TokenizerFast" if is_tokenizers_available() else None)),
+ ("qwen2_audio", ("Qwen2Tokenizer", "Qwen2TokenizerFast" if is_tokenizers_available() else None)),
+ (
+ "qwen2_moe",
+ (
+ "Qwen2Tokenizer",
+ "Qwen2TokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ ("qwen2_vl", ("Qwen2Tokenizer", "Qwen2TokenizerFast" if is_tokenizers_available() else None)),
+ (
+ "qwen3",
+ (
+ "Qwen2Tokenizer",
+ "Qwen2TokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ (
+ "qwen3_moe",
+ (
+ "Qwen2Tokenizer",
+ "Qwen2TokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ ("rag", ("RagTokenizer", None)),
+ ("realm", ("RealmTokenizer", "RealmTokenizerFast" if is_tokenizers_available() else None)),
+ (
+ "recurrent_gemma",
+ (
+ "GemmaTokenizer" if is_sentencepiece_available() else None,
+ "GemmaTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ (
+ "reformer",
+ (
+ "ReformerTokenizer" if is_sentencepiece_available() else None,
+ "ReformerTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ (
+ "rembert",
+ (
+ "RemBertTokenizer" if is_sentencepiece_available() else None,
+ "RemBertTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ ("retribert", ("RetriBertTokenizer", "RetriBertTokenizerFast" if is_tokenizers_available() else None)),
+ ("roberta", ("RobertaTokenizer", "RobertaTokenizerFast" if is_tokenizers_available() else None)),
+ (
+ "roberta-prelayernorm",
+ ("RobertaTokenizer", "RobertaTokenizerFast" if is_tokenizers_available() else None),
+ ),
+ ("roc_bert", ("RoCBertTokenizer", None)),
+ ("roformer", ("RoFormerTokenizer", "RoFormerTokenizerFast" if is_tokenizers_available() else None)),
+ ("rwkv", (None, "GPTNeoXTokenizerFast" if is_tokenizers_available() else None)),
+ (
+ "seamless_m4t",
+ (
+ "SeamlessM4TTokenizer" if is_sentencepiece_available() else None,
+ "SeamlessM4TTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ (
+ "seamless_m4t_v2",
+ (
+ "SeamlessM4TTokenizer" if is_sentencepiece_available() else None,
+ "SeamlessM4TTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ (
+ "shieldgemma2",
+ (
+ "GemmaTokenizer" if is_sentencepiece_available() else None,
+ "GemmaTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ ("siglip", ("SiglipTokenizer" if is_sentencepiece_available() else None, None)),
+ (
+ "siglip2",
+ (
+ "GemmaTokenizer" if is_sentencepiece_available() else None,
+ "GemmaTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ ("speech_to_text", ("Speech2TextTokenizer" if is_sentencepiece_available() else None, None)),
+ ("speech_to_text_2", ("Speech2Text2Tokenizer", None)),
+ ("speecht5", ("SpeechT5Tokenizer" if is_sentencepiece_available() else None, None)),
+ ("splinter", ("SplinterTokenizer", "SplinterTokenizerFast")),
+ (
+ "squeezebert",
+ ("SqueezeBertTokenizer", "SqueezeBertTokenizerFast" if is_tokenizers_available() else None),
+ ),
+ ("stablelm", (None, "GPTNeoXTokenizerFast" if is_tokenizers_available() else None)),
+ ("starcoder2", ("GPT2Tokenizer", "GPT2TokenizerFast" if is_tokenizers_available() else None)),
+ (
+ "switch_transformers",
+ (
+ "T5Tokenizer" if is_sentencepiece_available() else None,
+ "T5TokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ (
+ "t5",
+ (
+ "T5Tokenizer" if is_sentencepiece_available() else None,
+ "T5TokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ ("tapas", ("TapasTokenizer", None)),
+ ("tapex", ("TapexTokenizer", None)),
+ ("transfo-xl", ("TransfoXLTokenizer", None)),
+ ("tvp", ("BertTokenizer", "BertTokenizerFast" if is_tokenizers_available() else None)),
+ (
+ "udop",
+ (
+ "UdopTokenizer" if is_sentencepiece_available() else None,
+ "UdopTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ (
+ "umt5",
+ (
+ "T5Tokenizer" if is_sentencepiece_available() else None,
+ "T5TokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ ("video_llava", ("LlamaTokenizer", "LlamaTokenizerFast" if is_tokenizers_available() else None)),
+ ("vilt", ("BertTokenizer", "BertTokenizerFast" if is_tokenizers_available() else None)),
+ ("vipllava", ("LlamaTokenizer", "LlamaTokenizerFast" if is_tokenizers_available() else None)),
+ ("visual_bert", ("BertTokenizer", "BertTokenizerFast" if is_tokenizers_available() else None)),
+ ("vits", ("VitsTokenizer", None)),
+ ("wav2vec2", ("Wav2Vec2CTCTokenizer", None)),
+ ("wav2vec2-bert", ("Wav2Vec2CTCTokenizer", None)),
+ ("wav2vec2-conformer", ("Wav2Vec2CTCTokenizer", None)),
+ ("wav2vec2_phoneme", ("Wav2Vec2PhonemeCTCTokenizer", None)),
+ ("whisper", ("WhisperTokenizer", "WhisperTokenizerFast" if is_tokenizers_available() else None)),
+ ("xclip", ("CLIPTokenizer", "CLIPTokenizerFast" if is_tokenizers_available() else None)),
+ (
+ "xglm",
+ (
+ "XGLMTokenizer" if is_sentencepiece_available() else None,
+ "XGLMTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ ("xlm", ("XLMTokenizer", None)),
+ ("xlm-prophetnet", ("XLMProphetNetTokenizer" if is_sentencepiece_available() else None, None)),
+ (
+ "xlm-roberta",
+ (
+ "XLMRobertaTokenizer" if is_sentencepiece_available() else None,
+ "XLMRobertaTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ (
+ "xlm-roberta-xl",
+ (
+ "XLMRobertaTokenizer" if is_sentencepiece_available() else None,
+ "XLMRobertaTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ (
+ "xlnet",
+ (
+ "XLNetTokenizer" if is_sentencepiece_available() else None,
+ "XLNetTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ (
+ "xmod",
+ (
+ "XLMRobertaTokenizer" if is_sentencepiece_available() else None,
+ "XLMRobertaTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ (
+ "yoso",
+ (
+ "AlbertTokenizer" if is_sentencepiece_available() else None,
+ "AlbertTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ (
+ "zamba",
+ (
+ "LlamaTokenizer" if is_sentencepiece_available() else None,
+ "LlamaTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ (
+ "zamba2",
+ (
+ "LlamaTokenizer" if is_sentencepiece_available() else None,
+ "LlamaTokenizerFast" if is_tokenizers_available() else None,
+ ),
+ ),
+ ]
+ )
+
+TOKENIZER_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, TOKENIZER_MAPPING_NAMES)
+
+CONFIG_TO_TYPE = {v: k for k, v in CONFIG_MAPPING_NAMES.items()}
+
+
+def tokenizer_class_from_name(class_name: str):
+ if class_name == "PreTrainedTokenizerFast":
+ return PreTrainedTokenizerFast
+
+ for module_name, tokenizers in TOKENIZER_MAPPING_NAMES.items():
+ if class_name in tokenizers:
+ module_name = model_type_to_module_name(module_name)
+
+ module = importlib.import_module(f".{module_name}", "transformers.models")
+ try:
+ return getattr(module, class_name)
+ except AttributeError:
+ continue
+
+ for config, tokenizers in TOKENIZER_MAPPING._extra_content.items():
+ for tokenizer in tokenizers:
+ if getattr(tokenizer, "__name__", None) == class_name:
+ return tokenizer
+
+ # We did not fine the class, but maybe it's because a dep is missing. In that case, the class will be in the main
+ # init and we return the proper dummy to get an appropriate error message.
+ main_module = importlib.import_module("transformers")
+ if hasattr(main_module, class_name):
+ return getattr(main_module, class_name)
+
+ return None
+
+
+def get_tokenizer_config(
+ pretrained_model_name_or_path: Union[str, os.PathLike],
+ cache_dir: Optional[Union[str, os.PathLike]] = None,
+ force_download: bool = False,
+ resume_download: Optional[bool] = None,
+ proxies: Optional[Dict[str, str]] = None,
+ token: Optional[Union[bool, str]] = None,
+ revision: Optional[str] = None,
+ local_files_only: bool = False,
+ subfolder: str = "",
+ **kwargs,
+):
+ """
+ Loads the tokenizer configuration from a pretrained model tokenizer configuration.
+
+ Args:
+ pretrained_model_name_or_path (`str` or `os.PathLike`):
+ This can be either:
+
+ - a string, the *model id* of a pretrained model configuration hosted inside a model repo on
+ huggingface.co.
+ - a path to a *directory* containing a configuration file saved using the
+ [`~PreTrainedTokenizer.save_pretrained`] method, e.g., `./my_model_directory/`.
+
+ cache_dir (`str` or `os.PathLike`, *optional*):
+ Path to a directory in which a downloaded pretrained model configuration should be cached if the standard
+ cache should not be used.
+ force_download (`bool`, *optional*, defaults to `False`):
+ Whether or not to force to (re-)download the configuration files and override the cached versions if they
+ exist.
+ resume_download:
+ Deprecated and ignored. All downloads are now resumed by default when possible.
+ Will be removed in v5 of Transformers.
+ proxies (`Dict[str, str]`, *optional*):
+ A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',
+ 'http://hostname': 'foo.bar:4012'}.` The proxies are used on each request.
+ token (`str` or *bool*, *optional*):
+ The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated
+ when running `huggingface-cli login` (stored in `~/.huggingface`).
+ revision (`str`, *optional*, defaults to `"main"`):
+ The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
+ git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any
+ identifier allowed by git.
+ local_files_only (`bool`, *optional*, defaults to `False`):
+ If `True`, will only try to load the tokenizer configuration from local files.
+ subfolder (`str`, *optional*, defaults to `""`):
+ In case the tokenizer config is located inside a subfolder of the model repo on huggingface.co, you can
+ specify the folder name here.
+
+
+
+ Passing `token=True` is required when you want to use a private model.
+
+
+
+ Returns:
+ `Dict`: The configuration of the tokenizer.
+
+ Examples:
+
+ ```python
+ # Download configuration from huggingface.co and cache.
+ tokenizer_config = get_tokenizer_config("google-bert/bert-base-uncased")
+ # This model does not have a tokenizer config so the result will be an empty dict.
+ tokenizer_config = get_tokenizer_config("FacebookAI/xlm-roberta-base")
+
+ # Save a pretrained tokenizer locally and you can reload its config
+ from transformers import AutoTokenizer
+
+ tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-cased")
+ tokenizer.save_pretrained("tokenizer-test")
+ tokenizer_config = get_tokenizer_config("tokenizer-test")
+ ```"""
+ use_auth_token = kwargs.pop("use_auth_token", None)
+ if use_auth_token is not None:
+ warnings.warn(
+ "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.",
+ FutureWarning,
+ )
+ if token is not None:
+ raise ValueError("`token` and `use_auth_token` are both specified. Please set only the argument `token`.")
+ token = use_auth_token
+
+ commit_hash = kwargs.get("_commit_hash", None)
+ resolved_config_file = cached_file(
+ pretrained_model_name_or_path,
+ TOKENIZER_CONFIG_FILE,
+ cache_dir=cache_dir,
+ force_download=force_download,
+ resume_download=resume_download,
+ proxies=proxies,
+ token=token,
+ revision=revision,
+ local_files_only=local_files_only,
+ subfolder=subfolder,
+ _raise_exceptions_for_gated_repo=False,
+ _raise_exceptions_for_missing_entries=False,
+ _raise_exceptions_for_connection_errors=False,
+ _commit_hash=commit_hash,
+ )
+ if resolved_config_file is None:
+ logger.info("Could not locate the tokenizer configuration file, will try to use the model config instead.")
+ return {}
+ commit_hash = extract_commit_hash(resolved_config_file, commit_hash)
+
+ with open(resolved_config_file, encoding="utf-8") as reader:
+ result = json.load(reader)
+ result["_commit_hash"] = commit_hash
+ return result
+
+
+class AutoTokenizer:
+ r"""
+ This is a generic tokenizer class that will be instantiated as one of the tokenizer classes of the library when
+ created with the [`AutoTokenizer.from_pretrained`] class method.
+
+ This class cannot be instantiated directly using `__init__()` (throws an error).
+ """
+
+ def __init__(self):
+ raise EnvironmentError(
+ "AutoTokenizer is designed to be instantiated "
+ "using the `AutoTokenizer.from_pretrained(pretrained_model_name_or_path)` method."
+ )
+
+ @classmethod
+ @replace_list_option_in_docstrings(TOKENIZER_MAPPING_NAMES)
+ def from_pretrained(cls, pretrained_model_name_or_path, *inputs, **kwargs):
+ r"""
+ Instantiate one of the tokenizer classes of the library from a pretrained model vocabulary.
+
+ The tokenizer class to instantiate is selected based on the `model_type` property of the config object (either
+ passed as an argument or loaded from `pretrained_model_name_or_path` if possible), or when it's missing, by
+ falling back to using pattern matching on `pretrained_model_name_or_path`:
+
+ List options
+
+ Params:
+ pretrained_model_name_or_path (`str` or `os.PathLike`):
+ Can be either:
+
+ - A string, the *model id* of a predefined tokenizer hosted inside a model repo on huggingface.co.
+ - A path to a *directory* containing vocabulary files required by the tokenizer, for instance saved
+ using the [`~PreTrainedTokenizer.save_pretrained`] method, e.g., `./my_model_directory/`.
+ - A path or url to a single saved vocabulary file if and only if the tokenizer only requires a
+ single vocabulary file (like Bert or XLNet), e.g.: `./my_model_directory/vocab.txt`. (Not
+ applicable to all derived classes)
+ inputs (additional positional arguments, *optional*):
+ Will be passed along to the Tokenizer `__init__()` method.
+ config ([`PretrainedConfig`], *optional*)
+ The configuration object used to determine the tokenizer class to instantiate.
+ cache_dir (`str` or `os.PathLike`, *optional*):
+ Path to a directory in which a downloaded pretrained model configuration should be cached if the
+ standard cache should not be used.
+ force_download (`bool`, *optional*, defaults to `False`):
+ Whether or not to force the (re-)download the model weights and configuration files and override the
+ cached versions if they exist.
+ resume_download:
+ Deprecated and ignored. All downloads are now resumed by default when possible.
+ Will be removed in v5 of Transformers.
+ proxies (`Dict[str, str]`, *optional*):
+ A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',
+ 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
+ revision (`str`, *optional*, defaults to `"main"`):
+ The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
+ git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any
+ identifier allowed by git.
+ subfolder (`str`, *optional*):
+ In case the relevant files are located inside a subfolder of the model repo on huggingface.co (e.g. for
+ facebook/rag-token-base), specify it here.
+ use_fast (`bool`, *optional*, defaults to `True`):
+ Use a [fast Rust-based tokenizer](https://huggingface.co/docs/tokenizers/index) if it is supported for
+ a given model. If a fast tokenizer is not available for a given model, a normal Python-based tokenizer
+ is returned instead.
+ tokenizer_type (`str`, *optional*):
+ Tokenizer type to be loaded.
+ trust_remote_code (`bool`, *optional*, defaults to `False`):
+ Whether or not to allow for custom models defined on the Hub in their own modeling files. This option
+ should only be set to `True` for repositories you trust and in which you have read the code, as it will
+ execute code present on the Hub on your local machine.
+ kwargs (additional keyword arguments, *optional*):
+ Will be passed to the Tokenizer `__init__()` method. Can be used to set special tokens like
+ `bos_token`, `eos_token`, `unk_token`, `sep_token`, `pad_token`, `cls_token`, `mask_token`,
+ `additional_special_tokens`. See parameters in the `__init__()` for more details.
+
+ Examples:
+
+ ```python
+ >>> from transformers import AutoTokenizer
+
+ >>> # Download vocabulary from huggingface.co and cache.
+ >>> tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-uncased")
+
+ >>> # Download vocabulary from huggingface.co (user-uploaded) and cache.
+ >>> tokenizer = AutoTokenizer.from_pretrained("dbmdz/bert-base-german-cased")
+
+ >>> # If vocabulary files are in a directory (e.g. tokenizer was saved using *save_pretrained('./test/saved_model/')*)
+ >>> # tokenizer = AutoTokenizer.from_pretrained("./test/bert_saved_model/")
+
+ >>> # Download vocabulary from huggingface.co and define model-specific arguments
+ >>> tokenizer = AutoTokenizer.from_pretrained("FacebookAI/roberta-base", add_prefix_space=True)
+ ```"""
+ use_auth_token = kwargs.pop("use_auth_token", None)
+ if use_auth_token is not None:
+ warnings.warn(
+ "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.",
+ FutureWarning,
+ )
+ if kwargs.get("token", None) is not None:
+ raise ValueError(
+ "`token` and `use_auth_token` are both specified. Please set only the argument `token`."
+ )
+ kwargs["token"] = use_auth_token
+
+ config = kwargs.pop("config", None)
+ kwargs["_from_auto"] = True
+
+ use_fast = kwargs.pop("use_fast", True)
+ tokenizer_type = kwargs.pop("tokenizer_type", None)
+ trust_remote_code = kwargs.pop("trust_remote_code", None)
+ gguf_file = kwargs.get("gguf_file", None)
+
+ # First, let's see whether the tokenizer_type is passed so that we can leverage it
+ if tokenizer_type is not None:
+ tokenizer_class = None
+ tokenizer_class_tuple = TOKENIZER_MAPPING_NAMES.get(tokenizer_type, None)
+
+ if tokenizer_class_tuple is None:
+ raise ValueError(
+ f"Passed `tokenizer_type` {tokenizer_type} does not exist. `tokenizer_type` should be one of "
+ f"{', '.join(c for c in TOKENIZER_MAPPING_NAMES.keys())}."
+ )
+
+ tokenizer_class_name, tokenizer_fast_class_name = tokenizer_class_tuple
+
+ if use_fast:
+ if tokenizer_fast_class_name is not None:
+ tokenizer_class = tokenizer_class_from_name(tokenizer_fast_class_name)
+ else:
+ logger.warning(
+ "`use_fast` is set to `True` but the tokenizer class does not have a fast version. "
+ " Falling back to the slow version."
+ )
+ if tokenizer_class is None:
+ tokenizer_class = tokenizer_class_from_name(tokenizer_class_name)
+
+ if tokenizer_class is None:
+ raise ValueError(f"Tokenizer class {tokenizer_class_name} is not currently imported.")
+
+ return tokenizer_class.from_pretrained(pretrained_model_name_or_path, *inputs, **kwargs)
+
+ # Next, let's try to use the tokenizer_config file to get the tokenizer class.
+ tokenizer_config = get_tokenizer_config(pretrained_model_name_or_path, **kwargs)
+ if "_commit_hash" in tokenizer_config:
+ kwargs["_commit_hash"] = tokenizer_config["_commit_hash"]
+ config_tokenizer_class = tokenizer_config.get("tokenizer_class")
+ tokenizer_auto_map = None
+ if "auto_map" in tokenizer_config:
+ if isinstance(tokenizer_config["auto_map"], (tuple, list)):
+ # Legacy format for dynamic tokenizers
+ tokenizer_auto_map = tokenizer_config["auto_map"]
+ else:
+ tokenizer_auto_map = tokenizer_config["auto_map"].get("AutoTokenizer", None)
+
+ # If that did not work, let's try to use the config.
+ if config_tokenizer_class is None:
+ if not isinstance(config, PretrainedConfig):
+ if gguf_file:
+ gguf_path = cached_file(pretrained_model_name_or_path, gguf_file, **kwargs)
+ config_dict = load_gguf_checkpoint(gguf_path, return_tensors=False)["config"]
+ config = AutoConfig.for_model(**config_dict)
+ else:
+ config = AutoConfig.from_pretrained(
+ pretrained_model_name_or_path, trust_remote_code=trust_remote_code, **kwargs
+ )
+ config_tokenizer_class = config.tokenizer_class
+ if hasattr(config, "auto_map") and "AutoTokenizer" in config.auto_map:
+ tokenizer_auto_map = config.auto_map["AutoTokenizer"]
+
+ has_remote_code = tokenizer_auto_map is not None
+ has_local_code = type(config) in TOKENIZER_MAPPING or (
+ config_tokenizer_class is not None
+ and (
+ tokenizer_class_from_name(config_tokenizer_class) is not None
+ or tokenizer_class_from_name(config_tokenizer_class + "Fast") is not None
+ )
+ )
+ trust_remote_code = resolve_trust_remote_code(
+ trust_remote_code, pretrained_model_name_or_path, has_local_code, has_remote_code
+ )
+
+ if has_remote_code and trust_remote_code:
+ if use_fast and tokenizer_auto_map[1] is not None:
+ class_ref = tokenizer_auto_map[1]
+ else:
+ class_ref = tokenizer_auto_map[0]
+ tokenizer_class = get_class_from_dynamic_module(class_ref, pretrained_model_name_or_path, **kwargs)
+ _ = kwargs.pop("code_revision", None)
+ if os.path.isdir(pretrained_model_name_or_path):
+ tokenizer_class.register_for_auto_class()
+ return tokenizer_class.from_pretrained(
+ pretrained_model_name_or_path, *inputs, trust_remote_code=trust_remote_code, **kwargs
+ )
+ elif config_tokenizer_class is not None:
+ tokenizer_class = None
+ if use_fast and not config_tokenizer_class.endswith("Fast"):
+ tokenizer_class_candidate = f"{config_tokenizer_class}Fast"
+ tokenizer_class = tokenizer_class_from_name(tokenizer_class_candidate)
+ if tokenizer_class is None:
+ tokenizer_class_candidate = config_tokenizer_class
+ tokenizer_class = tokenizer_class_from_name(tokenizer_class_candidate)
+ if tokenizer_class is None:
+ raise ValueError(
+ f"Tokenizer class {tokenizer_class_candidate} does not exist or is not currently imported."
+ )
+ return tokenizer_class.from_pretrained(pretrained_model_name_or_path, *inputs, **kwargs)
+
+ # Otherwise we have to be creative.
+ # if model is an encoder decoder, the encoder tokenizer class is used by default
+ if isinstance(config, EncoderDecoderConfig):
+ if type(config.decoder) is not type(config.encoder): # noqa: E721
+ logger.warning(
+ f"The encoder model config class: {config.encoder.__class__} is different from the decoder model "
+ f"config class: {config.decoder.__class__}. It is not recommended to use the "
+ "`AutoTokenizer.from_pretrained()` method in this case. Please use the encoder and decoder "
+ "specific tokenizer classes."
+ )
+ config = config.encoder
+
+ model_type = config_class_to_model_type(type(config).__name__)
+ if model_type is not None:
+ tokenizer_class_py, tokenizer_class_fast = TOKENIZER_MAPPING[type(config)]
+
+ if tokenizer_class_fast and (use_fast or tokenizer_class_py is None):
+ return tokenizer_class_fast.from_pretrained(pretrained_model_name_or_path, *inputs, **kwargs)
+ else:
+ if tokenizer_class_py is not None:
+ return tokenizer_class_py.from_pretrained(pretrained_model_name_or_path, *inputs, **kwargs)
+ else:
+ raise ValueError(
+ "This tokenizer cannot be instantiated. Please make sure you have `sentencepiece` installed "
+ "in order to use this tokenizer."
+ )
+
+ raise ValueError(
+ f"Unrecognized configuration class {config.__class__} to build an AutoTokenizer.\n"
+ f"Model type should be one of {', '.join(c.__name__ for c in TOKENIZER_MAPPING.keys())}."
+ )
+
+ @staticmethod
+ def register(config_class, slow_tokenizer_class=None, fast_tokenizer_class=None, exist_ok=False):
+ """
+ Register a new tokenizer in this mapping.
+
+
+ Args:
+ config_class ([`PretrainedConfig`]):
+ The configuration corresponding to the model to register.
+ slow_tokenizer_class ([`PretrainedTokenizer`], *optional*):
+ The slow tokenizer to register.
+ fast_tokenizer_class ([`PretrainedTokenizerFast`], *optional*):
+ The fast tokenizer to register.
+ """
+ if slow_tokenizer_class is None and fast_tokenizer_class is None:
+ raise ValueError("You need to pass either a `slow_tokenizer_class` or a `fast_tokenizer_class")
+ if slow_tokenizer_class is not None and issubclass(slow_tokenizer_class, PreTrainedTokenizerFast):
+ raise ValueError("You passed a fast tokenizer in the `slow_tokenizer_class`.")
+ if fast_tokenizer_class is not None and issubclass(fast_tokenizer_class, PreTrainedTokenizer):
+ raise ValueError("You passed a slow tokenizer in the `fast_tokenizer_class`.")
+
+ if (
+ slow_tokenizer_class is not None
+ and fast_tokenizer_class is not None
+ and issubclass(fast_tokenizer_class, PreTrainedTokenizerFast)
+ and fast_tokenizer_class.slow_tokenizer_class != slow_tokenizer_class
+ ):
+ raise ValueError(
+ "The fast tokenizer class you are passing has a `slow_tokenizer_class` attribute that is not "
+ "consistent with the slow tokenizer class you passed (fast tokenizer has "
+ f"{fast_tokenizer_class.slow_tokenizer_class} and you passed {slow_tokenizer_class}. Fix one of those "
+ "so they match!"
+ )
+
+ # Avoid resetting a set slow/fast tokenizer if we are passing just the other ones.
+ if config_class in TOKENIZER_MAPPING._extra_content:
+ existing_slow, existing_fast = TOKENIZER_MAPPING[config_class]
+ if slow_tokenizer_class is None:
+ slow_tokenizer_class = existing_slow
+ if fast_tokenizer_class is None:
+ fast_tokenizer_class = existing_fast
+
+ TOKENIZER_MAPPING.register(config_class, (slow_tokenizer_class, fast_tokenizer_class), exist_ok=exist_ok)
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/tokenization_bert.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/tokenization_bert.py
new file mode 100644
index 0000000000000000000000000000000000000000..42d4dd94554d411fd37b6328b51887a47401f114
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/tokenization_bert.py
@@ -0,0 +1,507 @@
+# coding=utf-8
+# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. 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.
+"""Tokenization classes for Bert."""
+
+import collections
+import os
+import unicodedata
+from typing import List, Optional, Tuple
+
+from ...tokenization_utils import PreTrainedTokenizer, _is_control, _is_punctuation, _is_whitespace
+from ...utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt"}
+
+
+def load_vocab(vocab_file):
+ """Loads a vocabulary file into a dictionary."""
+ vocab = collections.OrderedDict()
+ with open(vocab_file, "r", encoding="utf-8") as reader:
+ tokens = reader.readlines()
+ for index, token in enumerate(tokens):
+ token = token.rstrip("\n")
+ vocab[token] = index
+ return vocab
+
+
+def whitespace_tokenize(text):
+ """Runs basic whitespace cleaning and splitting on a piece of text."""
+ text = text.strip()
+ if not text:
+ return []
+ tokens = text.split()
+ return tokens
+
+
+class BertTokenizer(PreTrainedTokenizer):
+ r"""
+ Construct a BERT tokenizer. Based on WordPiece.
+
+ This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to
+ this superclass for more information regarding those methods.
+
+ Args:
+ vocab_file (`str`):
+ File containing the vocabulary.
+ do_lower_case (`bool`, *optional*, defaults to `True`):
+ Whether or not to lowercase the input when tokenizing.
+ do_basic_tokenize (`bool`, *optional*, defaults to `True`):
+ Whether or not to do basic tokenization before WordPiece.
+ never_split (`Iterable`, *optional*):
+ Collection of tokens which will never be split during tokenization. Only has an effect when
+ `do_basic_tokenize=True`
+ unk_token (`str`, *optional*, defaults to `"[UNK]"`):
+ The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
+ token instead.
+ sep_token (`str`, *optional*, defaults to `"[SEP]"`):
+ The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for
+ sequence classification or for a text and a question for question answering. It is also used as the last
+ token of a sequence built with special tokens.
+ pad_token (`str`, *optional*, defaults to `"[PAD]"`):
+ The token used for padding, for example when batching sequences of different lengths.
+ cls_token (`str`, *optional*, defaults to `"[CLS]"`):
+ The classifier token which is used when doing sequence classification (classification of the whole sequence
+ instead of per-token classification). It is the first token of the sequence when built with special tokens.
+ mask_token (`str`, *optional*, defaults to `"[MASK]"`):
+ The token used for masking values. This is the token used when training this model with masked language
+ modeling. This is the token which the model will try to predict.
+ tokenize_chinese_chars (`bool`, *optional*, defaults to `True`):
+ Whether or not to tokenize Chinese characters.
+
+ This should likely be deactivated for Japanese (see this
+ [issue](https://github.com/huggingface/transformers/issues/328)).
+ strip_accents (`bool`, *optional*):
+ Whether or not to strip all accents. If this option is not specified, then it will be determined by the
+ value for `lowercase` (as in the original BERT).
+ clean_up_tokenization_spaces (`bool`, *optional*, defaults to `True`):
+ Whether or not to cleanup spaces after decoding, cleanup consists in removing potential artifacts like
+ extra spaces.
+ """
+
+ vocab_files_names = VOCAB_FILES_NAMES
+
+ def __init__(
+ self,
+ vocab_file,
+ do_lower_case=True,
+ do_basic_tokenize=True,
+ never_split=None,
+ unk_token="[UNK]",
+ sep_token="[SEP]",
+ pad_token="[PAD]",
+ cls_token="[CLS]",
+ mask_token="[MASK]",
+ tokenize_chinese_chars=True,
+ strip_accents=None,
+ clean_up_tokenization_spaces=True,
+ **kwargs,
+ ):
+ if not os.path.isfile(vocab_file):
+ raise ValueError(
+ f"Can't find a vocabulary file at path '{vocab_file}'. To load the vocabulary from a Google pretrained"
+ " model use `tokenizer = BertTokenizer.from_pretrained(PRETRAINED_MODEL_NAME)`"
+ )
+ self.vocab = load_vocab(vocab_file)
+ self.ids_to_tokens = collections.OrderedDict([(ids, tok) for tok, ids in self.vocab.items()])
+ self.do_basic_tokenize = do_basic_tokenize
+ if do_basic_tokenize:
+ self.basic_tokenizer = BasicTokenizer(
+ do_lower_case=do_lower_case,
+ never_split=never_split,
+ tokenize_chinese_chars=tokenize_chinese_chars,
+ strip_accents=strip_accents,
+ )
+
+ self.wordpiece_tokenizer = WordpieceTokenizer(vocab=self.vocab, unk_token=str(unk_token))
+
+ super().__init__(
+ do_lower_case=do_lower_case,
+ do_basic_tokenize=do_basic_tokenize,
+ never_split=never_split,
+ unk_token=unk_token,
+ sep_token=sep_token,
+ pad_token=pad_token,
+ cls_token=cls_token,
+ mask_token=mask_token,
+ tokenize_chinese_chars=tokenize_chinese_chars,
+ strip_accents=strip_accents,
+ clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+ **kwargs,
+ )
+
+ @property
+ def do_lower_case(self):
+ return self.basic_tokenizer.do_lower_case
+
+ @property
+ def vocab_size(self):
+ return len(self.vocab)
+
+ def get_vocab(self):
+ return dict(self.vocab, **self.added_tokens_encoder)
+
+ def _tokenize(self, text, split_special_tokens=False):
+ split_tokens = []
+ if self.do_basic_tokenize:
+ for token in self.basic_tokenizer.tokenize(
+ text, never_split=self.all_special_tokens if not split_special_tokens else None
+ ):
+ # If the token is part of the never_split set
+ if token in self.basic_tokenizer.never_split:
+ split_tokens.append(token)
+ else:
+ split_tokens += self.wordpiece_tokenizer.tokenize(token)
+ else:
+ split_tokens = self.wordpiece_tokenizer.tokenize(text)
+ return split_tokens
+
+ def _convert_token_to_id(self, token):
+ """Converts a token (str) in an id using the vocab."""
+ return self.vocab.get(token, self.vocab.get(self.unk_token))
+
+ def _convert_id_to_token(self, index):
+ """Converts an index (integer) in a token (str) using the vocab."""
+ return self.ids_to_tokens.get(index, self.unk_token)
+
+ def convert_tokens_to_string(self, tokens):
+ """Converts a sequence of tokens (string) in a single string."""
+ out_string = " ".join(tokens).replace(" ##", "").strip()
+ return out_string
+
+ def build_inputs_with_special_tokens(
+ self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
+ ) -> List[int]:
+ """
+ Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
+ adding special tokens. A BERT sequence has the following format:
+
+ - single sequence: `[CLS] X [SEP]`
+ - pair of sequences: `[CLS] A [SEP] B [SEP]`
+
+ Args:
+ token_ids_0 (`List[int]`):
+ List of IDs to which the special tokens will be added.
+ token_ids_1 (`List[int]`, *optional*):
+ Optional second list of IDs for sequence pairs.
+
+ Returns:
+ `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
+ """
+ if token_ids_1 is None:
+ return [self.cls_token_id] + token_ids_0 + [self.sep_token_id]
+ cls = [self.cls_token_id]
+ sep = [self.sep_token_id]
+ return cls + token_ids_0 + sep + token_ids_1 + sep
+
+ def get_special_tokens_mask(
+ self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False
+ ) -> List[int]:
+ """
+ Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
+ special tokens using the tokenizer `prepare_for_model` method.
+
+ Args:
+ token_ids_0 (`List[int]`):
+ List of IDs.
+ token_ids_1 (`List[int]`, *optional*):
+ Optional second list of IDs for sequence pairs.
+ already_has_special_tokens (`bool`, *optional*, defaults to `False`):
+ Whether or not the token list is already formatted with special tokens for the model.
+
+ Returns:
+ `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
+ """
+
+ if already_has_special_tokens:
+ return super().get_special_tokens_mask(
+ token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
+ )
+
+ if token_ids_1 is not None:
+ return [1] + ([0] * len(token_ids_0)) + [1] + ([0] * len(token_ids_1)) + [1]
+ return [1] + ([0] * len(token_ids_0)) + [1]
+
+ def create_token_type_ids_from_sequences(
+ self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
+ ) -> List[int]:
+ """
+ Create a mask from the two sequences passed to be used in a sequence-pair classification task. A BERT sequence
+ pair mask has the following format:
+
+ ```
+ 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1
+ | first sequence | second sequence |
+ ```
+
+ If `token_ids_1` is `None`, this method only returns the first portion of the mask (0s).
+
+ Args:
+ token_ids_0 (`List[int]`):
+ List of IDs.
+ token_ids_1 (`List[int]`, *optional*):
+ Optional second list of IDs for sequence pairs.
+
+ Returns:
+ `List[int]`: List of [token type IDs](../glossary#token-type-ids) according to the given sequence(s).
+ """
+ sep = [self.sep_token_id]
+ cls = [self.cls_token_id]
+ if token_ids_1 is None:
+ return len(cls + token_ids_0 + sep) * [0]
+ return len(cls + token_ids_0 + sep) * [0] + len(token_ids_1 + sep) * [1]
+
+ def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
+ index = 0
+ if os.path.isdir(save_directory):
+ vocab_file = os.path.join(
+ save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
+ )
+ else:
+ vocab_file = (filename_prefix + "-" if filename_prefix else "") + save_directory
+ with open(vocab_file, "w", encoding="utf-8") as writer:
+ for token, token_index in sorted(self.vocab.items(), key=lambda kv: kv[1]):
+ if index != token_index:
+ logger.warning(
+ f"Saving vocabulary to {vocab_file}: vocabulary indices are not consecutive."
+ " Please check that the vocabulary is not corrupted!"
+ )
+ index = token_index
+ writer.write(token + "\n")
+ index += 1
+ return (vocab_file,)
+
+
+class BasicTokenizer:
+ """
+ Constructs a BasicTokenizer that will run basic tokenization (punctuation splitting, lower casing, etc.).
+
+ Args:
+ do_lower_case (`bool`, *optional*, defaults to `True`):
+ Whether or not to lowercase the input when tokenizing.
+ never_split (`Iterable`, *optional*):
+ Collection of tokens which will never be split during tokenization. Only has an effect when
+ `do_basic_tokenize=True`
+ tokenize_chinese_chars (`bool`, *optional*, defaults to `True`):
+ Whether or not to tokenize Chinese characters.
+
+ This should likely be deactivated for Japanese (see this
+ [issue](https://github.com/huggingface/transformers/issues/328)).
+ strip_accents (`bool`, *optional*):
+ Whether or not to strip all accents. If this option is not specified, then it will be determined by the
+ value for `lowercase` (as in the original BERT).
+ do_split_on_punc (`bool`, *optional*, defaults to `True`):
+ In some instances we want to skip the basic punctuation splitting so that later tokenization can capture
+ the full context of the words, such as contractions.
+ """
+
+ def __init__(
+ self,
+ do_lower_case=True,
+ never_split=None,
+ tokenize_chinese_chars=True,
+ strip_accents=None,
+ do_split_on_punc=True,
+ ):
+ if never_split is None:
+ never_split = []
+ self.do_lower_case = do_lower_case
+ self.never_split = set(never_split)
+ self.tokenize_chinese_chars = tokenize_chinese_chars
+ self.strip_accents = strip_accents
+ self.do_split_on_punc = do_split_on_punc
+
+ def tokenize(self, text, never_split=None):
+ """
+ Basic Tokenization of a piece of text. For sub-word tokenization, see WordPieceTokenizer.
+
+ Args:
+ never_split (`List[str]`, *optional*)
+ Kept for backward compatibility purposes. Now implemented directly at the base class level (see
+ [`PreTrainedTokenizer.tokenize`]) List of token not to split.
+ """
+ # union() returns a new set by concatenating the two sets.
+ never_split = self.never_split.union(set(never_split)) if never_split else self.never_split
+ text = self._clean_text(text)
+
+ # This was added on November 1st, 2018 for the multilingual and Chinese
+ # models. This is also applied to the English models now, but it doesn't
+ # matter since the English models were not trained on any Chinese data
+ # and generally don't have any Chinese data in them (there are Chinese
+ # characters in the vocabulary because Wikipedia does have some Chinese
+ # words in the English Wikipedia.).
+ if self.tokenize_chinese_chars:
+ text = self._tokenize_chinese_chars(text)
+ # prevents treating the same character with different unicode codepoints as different characters
+ unicode_normalized_text = unicodedata.normalize("NFC", text)
+ orig_tokens = whitespace_tokenize(unicode_normalized_text)
+ split_tokens = []
+ for token in orig_tokens:
+ if token not in never_split:
+ if self.do_lower_case:
+ token = token.lower()
+ if self.strip_accents is not False:
+ token = self._run_strip_accents(token)
+ elif self.strip_accents:
+ token = self._run_strip_accents(token)
+ split_tokens.extend(self._run_split_on_punc(token, never_split))
+
+ output_tokens = whitespace_tokenize(" ".join(split_tokens))
+ return output_tokens
+
+ def _run_strip_accents(self, text):
+ """Strips accents from a piece of text."""
+ text = unicodedata.normalize("NFD", text)
+ output = []
+ for char in text:
+ cat = unicodedata.category(char)
+ if cat == "Mn":
+ continue
+ output.append(char)
+ return "".join(output)
+
+ def _run_split_on_punc(self, text, never_split=None):
+ """Splits punctuation on a piece of text."""
+ if not self.do_split_on_punc or (never_split is not None and text in never_split):
+ return [text]
+ chars = list(text)
+ i = 0
+ start_new_word = True
+ output = []
+ while i < len(chars):
+ char = chars[i]
+ if _is_punctuation(char):
+ output.append([char])
+ start_new_word = True
+ else:
+ if start_new_word:
+ output.append([])
+ start_new_word = False
+ output[-1].append(char)
+ i += 1
+
+ return ["".join(x) for x in output]
+
+ def _tokenize_chinese_chars(self, text):
+ """Adds whitespace around any CJK character."""
+ output = []
+ for char in text:
+ cp = ord(char)
+ if self._is_chinese_char(cp):
+ output.append(" ")
+ output.append(char)
+ output.append(" ")
+ else:
+ output.append(char)
+ return "".join(output)
+
+ def _is_chinese_char(self, cp):
+ """Checks whether CP is the codepoint of a CJK character."""
+ # This defines a "chinese character" as anything in the CJK Unicode block:
+ # https://en.wikipedia.org/wiki/CJK_Unified_Ideographs_(Unicode_block)
+ #
+ # Note that the CJK Unicode block is NOT all Japanese and Korean characters,
+ # despite its name. The modern Korean Hangul alphabet is a different block,
+ # as is Japanese Hiragana and Katakana. Those alphabets are used to write
+ # space-separated words, so they are not treated specially and handled
+ # like the all of the other languages.
+ if (
+ (cp >= 0x4E00 and cp <= 0x9FFF)
+ or (cp >= 0x3400 and cp <= 0x4DBF) #
+ or (cp >= 0x20000 and cp <= 0x2A6DF) #
+ or (cp >= 0x2A700 and cp <= 0x2B73F) #
+ or (cp >= 0x2B740 and cp <= 0x2B81F) #
+ or (cp >= 0x2B820 and cp <= 0x2CEAF) #
+ or (cp >= 0xF900 and cp <= 0xFAFF)
+ or (cp >= 0x2F800 and cp <= 0x2FA1F) #
+ ): #
+ return True
+
+ return False
+
+ def _clean_text(self, text):
+ """Performs invalid character removal and whitespace cleanup on text."""
+ output = []
+ for char in text:
+ cp = ord(char)
+ if cp == 0 or cp == 0xFFFD or _is_control(char):
+ continue
+ if _is_whitespace(char):
+ output.append(" ")
+ else:
+ output.append(char)
+ return "".join(output)
+
+
+class WordpieceTokenizer:
+ """Runs WordPiece tokenization."""
+
+ def __init__(self, vocab, unk_token, max_input_chars_per_word=100):
+ self.vocab = vocab
+ self.unk_token = unk_token
+ self.max_input_chars_per_word = max_input_chars_per_word
+
+ def tokenize(self, text):
+ """
+ Tokenizes a piece of text into its word pieces. This uses a greedy longest-match-first algorithm to perform
+ tokenization using the given vocabulary.
+
+ For example, `input = "unaffable"` wil return as output `["un", "##aff", "##able"]`.
+
+ Args:
+ text: A single token or whitespace separated tokens. This should have
+ already been passed through *BasicTokenizer*.
+
+ Returns:
+ A list of wordpiece tokens.
+ """
+
+ output_tokens = []
+ for token in whitespace_tokenize(text):
+ chars = list(token)
+ if len(chars) > self.max_input_chars_per_word:
+ output_tokens.append(self.unk_token)
+ continue
+
+ is_bad = False
+ start = 0
+ sub_tokens = []
+ while start < len(chars):
+ end = len(chars)
+ cur_substr = None
+ while start < end:
+ substr = "".join(chars[start:end])
+ if start > 0:
+ substr = "##" + substr
+ if substr in self.vocab:
+ cur_substr = substr
+ break
+ end -= 1
+ if cur_substr is None:
+ is_bad = True
+ break
+ sub_tokens.append(cur_substr)
+ start = end
+
+ if is_bad:
+ output_tokens.append(self.unk_token)
+ else:
+ output_tokens.extend(sub_tokens)
+ return output_tokens
+
+
+__all__ = ["BasicTokenizer", "BertTokenizer", "WordpieceTokenizer"]
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/tokenization_bert_fast.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/tokenization_bert_fast.py
new file mode 100644
index 0000000000000000000000000000000000000000..4a89e6053b988f5b9d8131304c7b3e7e74dba5fd
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/tokenization_bert_fast.py
@@ -0,0 +1,175 @@
+# coding=utf-8
+# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. 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.
+"""Fast Tokenization classes for Bert."""
+
+import json
+from typing import List, Optional, Tuple
+
+from tokenizers import normalizers
+
+from ...tokenization_utils_fast import PreTrainedTokenizerFast
+from ...utils import logging
+from .tokenization_bert import BertTokenizer
+
+
+logger = logging.get_logger(__name__)
+
+VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt", "tokenizer_file": "tokenizer.json"}
+
+
+class BertTokenizerFast(PreTrainedTokenizerFast):
+ r"""
+ Construct a "fast" BERT tokenizer (backed by HuggingFace's *tokenizers* library). Based on WordPiece.
+
+ This tokenizer inherits from [`PreTrainedTokenizerFast`] which contains most of the main methods. Users should
+ refer to this superclass for more information regarding those methods.
+
+ Args:
+ vocab_file (`str`):
+ File containing the vocabulary.
+ do_lower_case (`bool`, *optional*, defaults to `True`):
+ Whether or not to lowercase the input when tokenizing.
+ unk_token (`str`, *optional*, defaults to `"[UNK]"`):
+ The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
+ token instead.
+ sep_token (`str`, *optional*, defaults to `"[SEP]"`):
+ The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for
+ sequence classification or for a text and a question for question answering. It is also used as the last
+ token of a sequence built with special tokens.
+ pad_token (`str`, *optional*, defaults to `"[PAD]"`):
+ The token used for padding, for example when batching sequences of different lengths.
+ cls_token (`str`, *optional*, defaults to `"[CLS]"`):
+ The classifier token which is used when doing sequence classification (classification of the whole sequence
+ instead of per-token classification). It is the first token of the sequence when built with special tokens.
+ mask_token (`str`, *optional*, defaults to `"[MASK]"`):
+ The token used for masking values. This is the token used when training this model with masked language
+ modeling. This is the token which the model will try to predict.
+ clean_text (`bool`, *optional*, defaults to `True`):
+ Whether or not to clean the text before tokenization by removing any control characters and replacing all
+ whitespaces by the classic one.
+ tokenize_chinese_chars (`bool`, *optional*, defaults to `True`):
+ Whether or not to tokenize Chinese characters. This should likely be deactivated for Japanese (see [this
+ issue](https://github.com/huggingface/transformers/issues/328)).
+ strip_accents (`bool`, *optional*):
+ Whether or not to strip all accents. If this option is not specified, then it will be determined by the
+ value for `lowercase` (as in the original BERT).
+ wordpieces_prefix (`str`, *optional*, defaults to `"##"`):
+ The prefix for subwords.
+ """
+
+ vocab_files_names = VOCAB_FILES_NAMES
+ slow_tokenizer_class = BertTokenizer
+
+ def __init__(
+ self,
+ vocab_file=None,
+ tokenizer_file=None,
+ do_lower_case=True,
+ unk_token="[UNK]",
+ sep_token="[SEP]",
+ pad_token="[PAD]",
+ cls_token="[CLS]",
+ mask_token="[MASK]",
+ tokenize_chinese_chars=True,
+ strip_accents=None,
+ **kwargs,
+ ):
+ super().__init__(
+ vocab_file,
+ tokenizer_file=tokenizer_file,
+ do_lower_case=do_lower_case,
+ unk_token=unk_token,
+ sep_token=sep_token,
+ pad_token=pad_token,
+ cls_token=cls_token,
+ mask_token=mask_token,
+ tokenize_chinese_chars=tokenize_chinese_chars,
+ strip_accents=strip_accents,
+ **kwargs,
+ )
+
+ normalizer_state = json.loads(self.backend_tokenizer.normalizer.__getstate__())
+ if (
+ normalizer_state.get("lowercase", do_lower_case) != do_lower_case
+ or normalizer_state.get("strip_accents", strip_accents) != strip_accents
+ or normalizer_state.get("handle_chinese_chars", tokenize_chinese_chars) != tokenize_chinese_chars
+ ):
+ normalizer_class = getattr(normalizers, normalizer_state.pop("type"))
+ normalizer_state["lowercase"] = do_lower_case
+ normalizer_state["strip_accents"] = strip_accents
+ normalizer_state["handle_chinese_chars"] = tokenize_chinese_chars
+ self.backend_tokenizer.normalizer = normalizer_class(**normalizer_state)
+
+ self.do_lower_case = do_lower_case
+
+ def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
+ """
+ Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
+ adding special tokens. A BERT sequence has the following format:
+
+ - single sequence: `[CLS] X [SEP]`
+ - pair of sequences: `[CLS] A [SEP] B [SEP]`
+
+ Args:
+ token_ids_0 (`List[int]`):
+ List of IDs to which the special tokens will be added.
+ token_ids_1 (`List[int]`, *optional*):
+ Optional second list of IDs for sequence pairs.
+
+ Returns:
+ `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
+ """
+ output = [self.cls_token_id] + token_ids_0 + [self.sep_token_id]
+
+ if token_ids_1 is not None:
+ output += token_ids_1 + [self.sep_token_id]
+
+ return output
+
+ def create_token_type_ids_from_sequences(
+ self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
+ ) -> List[int]:
+ """
+ Create a mask from the two sequences passed to be used in a sequence-pair classification task. A BERT sequence
+ pair mask has the following format:
+
+ ```
+ 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1
+ | first sequence | second sequence |
+ ```
+
+ If `token_ids_1` is `None`, this method only returns the first portion of the mask (0s).
+
+ Args:
+ token_ids_0 (`List[int]`):
+ List of IDs.
+ token_ids_1 (`List[int]`, *optional*):
+ Optional second list of IDs for sequence pairs.
+
+ Returns:
+ `List[int]`: List of [token type IDs](../glossary#token-type-ids) according to the given sequence(s).
+ """
+ sep = [self.sep_token_id]
+ cls = [self.cls_token_id]
+ if token_ids_1 is None:
+ return len(cls + token_ids_0 + sep) * [0]
+ return len(cls + token_ids_0 + sep) * [0] + len(token_ids_1 + sep) * [1]
+
+ def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
+ files = self._tokenizer.model.save(save_directory, name=filename_prefix)
+ return tuple(files)
+
+
+__all__ = ["BertTokenizerFast"]
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/tokenization_bert_tf.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/tokenization_bert_tf.py
new file mode 100644
index 0000000000000000000000000000000000000000..86658de5241a7ad89f44c6166bc8ddbba20f9502
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/tokenization_bert_tf.py
@@ -0,0 +1,257 @@
+import os
+from typing import List, Optional, Union
+
+import tensorflow as tf
+from tensorflow_text import BertTokenizer as BertTokenizerLayer
+from tensorflow_text import FastBertTokenizer, ShrinkLongestTrimmer, case_fold_utf8, combine_segments, pad_model_inputs
+
+from ...modeling_tf_utils import keras
+from .tokenization_bert import BertTokenizer
+
+
+class TFBertTokenizer(keras.layers.Layer):
+ """
+ This is an in-graph tokenizer for BERT. It should be initialized similarly to other tokenizers, using the
+ `from_pretrained()` method. It can also be initialized with the `from_tokenizer()` method, which imports settings
+ from an existing standard tokenizer object.
+
+ In-graph tokenizers, unlike other Hugging Face tokenizers, are actually Keras layers and are designed to be run
+ when the model is called, rather than during preprocessing. As a result, they have somewhat more limited options
+ than standard tokenizer classes. They are most useful when you want to create an end-to-end model that goes
+ straight from `tf.string` inputs to outputs.
+
+ Args:
+ vocab_list (`list`):
+ List containing the vocabulary.
+ do_lower_case (`bool`, *optional*, defaults to `True`):
+ Whether or not to lowercase the input when tokenizing.
+ cls_token_id (`str`, *optional*, defaults to `"[CLS]"`):
+ The classifier token which is used when doing sequence classification (classification of the whole sequence
+ instead of per-token classification). It is the first token of the sequence when built with special tokens.
+ sep_token_id (`str`, *optional*, defaults to `"[SEP]"`):
+ The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for
+ sequence classification or for a text and a question for question answering. It is also used as the last
+ token of a sequence built with special tokens.
+ pad_token_id (`str`, *optional*, defaults to `"[PAD]"`):
+ The token used for padding, for example when batching sequences of different lengths.
+ padding (`str`, defaults to `"longest"`):
+ The type of padding to use. Can be either `"longest"`, to pad only up to the longest sample in the batch,
+ or `"max_length", to pad all inputs to the maximum length supported by the tokenizer.
+ truncation (`bool`, *optional*, defaults to `True`):
+ Whether to truncate the sequence to the maximum length.
+ max_length (`int`, *optional*, defaults to `512`):
+ The maximum length of the sequence, used for padding (if `padding` is "max_length") and/or truncation (if
+ `truncation` is `True`).
+ pad_to_multiple_of (`int`, *optional*, defaults to `None`):
+ If set, the sequence will be padded to a multiple of this value.
+ return_token_type_ids (`bool`, *optional*, defaults to `True`):
+ Whether to return token_type_ids.
+ return_attention_mask (`bool`, *optional*, defaults to `True`):
+ Whether to return the attention_mask.
+ use_fast_bert_tokenizer (`bool`, *optional*, defaults to `True`):
+ If True, will use the FastBertTokenizer class from Tensorflow Text. If False, will use the BertTokenizer
+ class instead. BertTokenizer supports some additional options, but is slower and cannot be exported to
+ TFLite.
+ """
+
+ def __init__(
+ self,
+ vocab_list: List,
+ do_lower_case: bool,
+ cls_token_id: Optional[int] = None,
+ sep_token_id: Optional[int] = None,
+ pad_token_id: Optional[int] = None,
+ padding: str = "longest",
+ truncation: bool = True,
+ max_length: int = 512,
+ pad_to_multiple_of: Optional[int] = None,
+ return_token_type_ids: bool = True,
+ return_attention_mask: bool = True,
+ use_fast_bert_tokenizer: bool = True,
+ **tokenizer_kwargs,
+ ):
+ super().__init__()
+ if use_fast_bert_tokenizer:
+ self.tf_tokenizer = FastBertTokenizer(
+ vocab_list, token_out_type=tf.int64, lower_case_nfd_strip_accents=do_lower_case, **tokenizer_kwargs
+ )
+ else:
+ lookup_table = tf.lookup.StaticVocabularyTable(
+ tf.lookup.KeyValueTensorInitializer(
+ keys=vocab_list,
+ key_dtype=tf.string,
+ values=tf.range(tf.size(vocab_list, out_type=tf.int64), dtype=tf.int64),
+ value_dtype=tf.int64,
+ ),
+ num_oov_buckets=1,
+ )
+ self.tf_tokenizer = BertTokenizerLayer(
+ lookup_table, token_out_type=tf.int64, lower_case=do_lower_case, **tokenizer_kwargs
+ )
+
+ self.vocab_list = vocab_list
+ self.do_lower_case = do_lower_case
+ self.cls_token_id = vocab_list.index("[CLS]") if cls_token_id is None else cls_token_id
+ self.sep_token_id = vocab_list.index("[SEP]") if sep_token_id is None else sep_token_id
+ self.pad_token_id = vocab_list.index("[PAD]") if pad_token_id is None else pad_token_id
+ self.paired_trimmer = ShrinkLongestTrimmer(max_length - 3, axis=1) # Allow room for special tokens
+ self.max_length = max_length
+ self.padding = padding
+ self.truncation = truncation
+ self.pad_to_multiple_of = pad_to_multiple_of
+ self.return_token_type_ids = return_token_type_ids
+ self.return_attention_mask = return_attention_mask
+
+ @classmethod
+ def from_tokenizer(cls, tokenizer: "PreTrainedTokenizerBase", **kwargs): # noqa: F821
+ """
+ Initialize a `TFBertTokenizer` from an existing `Tokenizer`.
+
+ Args:
+ tokenizer (`PreTrainedTokenizerBase`):
+ The tokenizer to use to initialize the `TFBertTokenizer`.
+
+ Examples:
+
+ ```python
+ from transformers import AutoTokenizer, TFBertTokenizer
+
+ tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-uncased")
+ tf_tokenizer = TFBertTokenizer.from_tokenizer(tokenizer)
+ ```
+ """
+ do_lower_case = kwargs.pop("do_lower_case", None)
+ do_lower_case = tokenizer.do_lower_case if do_lower_case is None else do_lower_case
+ cls_token_id = kwargs.pop("cls_token_id", None)
+ cls_token_id = tokenizer.cls_token_id if cls_token_id is None else cls_token_id
+ sep_token_id = kwargs.pop("sep_token_id", None)
+ sep_token_id = tokenizer.sep_token_id if sep_token_id is None else sep_token_id
+ pad_token_id = kwargs.pop("pad_token_id", None)
+ pad_token_id = tokenizer.pad_token_id if pad_token_id is None else pad_token_id
+
+ vocab = tokenizer.get_vocab()
+ vocab = sorted(vocab.items(), key=lambda x: x[1])
+ vocab_list = [entry[0] for entry in vocab]
+ return cls(
+ vocab_list=vocab_list,
+ do_lower_case=do_lower_case,
+ cls_token_id=cls_token_id,
+ sep_token_id=sep_token_id,
+ pad_token_id=pad_token_id,
+ **kwargs,
+ )
+
+ @classmethod
+ def from_pretrained(cls, pretrained_model_name_or_path: Union[str, os.PathLike], *init_inputs, **kwargs):
+ """
+ Instantiate a `TFBertTokenizer` from a pre-trained tokenizer.
+
+ Args:
+ pretrained_model_name_or_path (`str` or `os.PathLike`):
+ The name or path to the pre-trained tokenizer.
+
+ Examples:
+
+ ```python
+ from transformers import TFBertTokenizer
+
+ tf_tokenizer = TFBertTokenizer.from_pretrained("google-bert/bert-base-uncased")
+ ```
+ """
+ try:
+ tokenizer = BertTokenizer.from_pretrained(pretrained_model_name_or_path, *init_inputs, **kwargs)
+ except: # noqa: E722
+ from .tokenization_bert_fast import BertTokenizerFast
+
+ tokenizer = BertTokenizerFast.from_pretrained(pretrained_model_name_or_path, *init_inputs, **kwargs)
+ return cls.from_tokenizer(tokenizer, **kwargs)
+
+ def unpaired_tokenize(self, texts):
+ if self.do_lower_case:
+ texts = case_fold_utf8(texts)
+ tokens = self.tf_tokenizer.tokenize(texts)
+ return tokens.merge_dims(1, -1)
+
+ def call(
+ self,
+ text,
+ text_pair=None,
+ padding=None,
+ truncation=None,
+ max_length=None,
+ pad_to_multiple_of=None,
+ return_token_type_ids=None,
+ return_attention_mask=None,
+ ):
+ if padding is None:
+ padding = self.padding
+ if padding not in ("longest", "max_length"):
+ raise ValueError("Padding must be either 'longest' or 'max_length'!")
+ if max_length is not None and text_pair is not None:
+ # Because we have to instantiate a Trimmer to do it properly
+ raise ValueError("max_length cannot be overridden at call time when truncating paired texts!")
+ if max_length is None:
+ max_length = self.max_length
+ if truncation is None:
+ truncation = self.truncation
+ if pad_to_multiple_of is None:
+ pad_to_multiple_of = self.pad_to_multiple_of
+ if return_token_type_ids is None:
+ return_token_type_ids = self.return_token_type_ids
+ if return_attention_mask is None:
+ return_attention_mask = self.return_attention_mask
+ if not isinstance(text, tf.Tensor):
+ text = tf.convert_to_tensor(text)
+ if text_pair is not None and not isinstance(text_pair, tf.Tensor):
+ text_pair = tf.convert_to_tensor(text_pair)
+ if text_pair is not None:
+ if text.shape.rank > 1:
+ raise ValueError("text argument should not be multidimensional when a text pair is supplied!")
+ if text_pair.shape.rank > 1:
+ raise ValueError("text_pair should not be multidimensional!")
+ if text.shape.rank == 2:
+ text, text_pair = text[:, 0], text[:, 1]
+ text = self.unpaired_tokenize(text)
+ if text_pair is None: # Unpaired text
+ if truncation:
+ text = text[:, : max_length - 2] # Allow room for special tokens
+ input_ids, token_type_ids = combine_segments(
+ (text,), start_of_sequence_id=self.cls_token_id, end_of_segment_id=self.sep_token_id
+ )
+ else: # Paired text
+ text_pair = self.unpaired_tokenize(text_pair)
+ if truncation:
+ text, text_pair = self.paired_trimmer.trim([text, text_pair])
+ input_ids, token_type_ids = combine_segments(
+ (text, text_pair), start_of_sequence_id=self.cls_token_id, end_of_segment_id=self.sep_token_id
+ )
+ if padding == "longest":
+ pad_length = input_ids.bounding_shape(axis=1)
+ if pad_to_multiple_of is not None:
+ # No ceiling division in tensorflow, so we negate floordiv instead
+ pad_length = pad_to_multiple_of * (-tf.math.floordiv(-pad_length, pad_to_multiple_of))
+ else:
+ pad_length = max_length
+
+ input_ids, attention_mask = pad_model_inputs(input_ids, max_seq_length=pad_length, pad_value=self.pad_token_id)
+ output = {"input_ids": input_ids}
+ if return_attention_mask:
+ output["attention_mask"] = attention_mask
+ if return_token_type_ids:
+ token_type_ids, _ = pad_model_inputs(
+ token_type_ids, max_seq_length=pad_length, pad_value=self.pad_token_id
+ )
+ output["token_type_ids"] = token_type_ids
+ return output
+
+ def get_config(self):
+ return {
+ "vocab_list": self.vocab_list,
+ "do_lower_case": self.do_lower_case,
+ "cls_token_id": self.cls_token_id,
+ "sep_token_id": self.sep_token_id,
+ "pad_token_id": self.pad_token_id,
+ }
+
+
+__all__ = ["TFBertTokenizer"]
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/tokenization_utils.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/tokenization_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..4c9cf6cb0a0b7f08ba43e8761b02862a5c368951
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/tokenization_utils.py
@@ -0,0 +1,68 @@
+# coding=utf-8
+# Copyright 2021 The HuggingFace Inc. 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.
+"""Tokenization utils for RoFormer."""
+
+from typing import List
+
+from tokenizers import NormalizedString, PreTokenizedString, normalizers
+
+
+class JiebaPreTokenizer:
+ def __init__(self, vocab) -> None:
+ self.vocab = vocab
+ self.normalizers = normalizers.BertNormalizer(
+ clean_text=False,
+ handle_chinese_chars=True,
+ strip_accents=False,
+ lowercase=False,
+ )
+ try:
+ import rjieba
+ except ImportError:
+ raise ImportError(
+ "You need to install rjieba to use RoFormerTokenizer. "
+ "See https://pypi.org/project/rjieba/ for installation."
+ )
+ self.jieba = rjieba
+
+ def jieba_split(self, i: int, normalized_string: NormalizedString) -> List[NormalizedString]:
+ splits = []
+
+ # this code slice normalized_string is too slow (6s) but test_alignment_methods can pass
+ for token, start, end in self.jieba.tokenize(str(normalized_string), hmm=False):
+ if token in self.vocab:
+ splits.append(normalized_string[start:end])
+ else:
+ token_list = self.normalizers.normalize_str(token).split()
+ for token in token_list:
+ if token:
+ end = start + len(token)
+ splits.append(normalized_string[start:end])
+ start = end
+
+ # this code test_alignment_methods can't pass but fast (300ms)
+ # for token in self.jieba.cut(str(normalized_string), False):
+ # if token in self.vocab:
+ # splits.append(NormalizedString(token))
+ # else:
+ # token_list = self.normalizers.normalize_str(token).split()
+ # for token in token_list:
+ # if token:
+ # splits.append(NormalizedString(token))
+
+ return splits
+
+ def pre_tokenize(self, pretok: PreTokenizedString):
+ pretok.split(self.jieba_split)
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/tokenization_utils_base.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/tokenization_utils_base.py
new file mode 100644
index 0000000000000000000000000000000000000000..8e047569e7ce56573949b4e03aa9f1057802bf90
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/tokenization_utils_base.py
@@ -0,0 +1,4176 @@
+# coding=utf-8
+# Copyright 2020 The HuggingFace Inc. 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.
+"""
+Base classes common to both the slow and the fast tokenization classes: PreTrainedTokenizerBase (host all the user
+fronting encoding methods) Special token mixing (host the special tokens logic) and BatchEncoding (wrap the dictionary
+of output with special method for the Fast tokenizers)
+"""
+
+import copy
+import json
+import os
+import re
+import warnings
+from collections import UserDict
+from collections.abc import Mapping, Sized
+from contextlib import contextmanager
+from dataclasses import dataclass
+from inspect import isfunction
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, NamedTuple, Optional, Sequence, Tuple, Union
+
+import numpy as np
+from packaging import version
+
+from . import __version__
+from .dynamic_module_utils import custom_object_save
+from .utils import (
+ ExplicitEnum,
+ PaddingStrategy,
+ PushToHubMixin,
+ TensorType,
+ add_end_docstrings,
+ add_model_info_to_auto_map,
+ add_model_info_to_custom_pipelines,
+ cached_file,
+ copy_func,
+ download_url,
+ extract_commit_hash,
+ get_json_schema,
+ is_flax_available,
+ is_jax_tensor,
+ is_mlx_available,
+ is_numpy_array,
+ is_offline_mode,
+ is_protobuf_available,
+ is_remote_url,
+ is_tf_available,
+ is_tf_tensor,
+ is_tokenizers_available,
+ is_torch_available,
+ is_torch_device,
+ is_torch_tensor,
+ logging,
+ requires_backends,
+ to_py_obj,
+)
+from .utils.chat_template_utils import _compile_jinja_template, _render_with_assistant_indices
+from .utils.import_utils import PROTOBUF_IMPORT_ERROR
+
+
+if TYPE_CHECKING:
+ if is_torch_available():
+ import torch
+ if is_tf_available():
+ import tensorflow as tf
+ if is_flax_available():
+ import jax.numpy as jnp # noqa: F401
+
+
+def import_protobuf_decode_error(error_message=""):
+ if is_protobuf_available():
+ from google.protobuf.message import DecodeError
+
+ return DecodeError
+ else:
+ raise ImportError(PROTOBUF_IMPORT_ERROR.format(error_message))
+
+
+if is_tokenizers_available():
+ from tokenizers import AddedToken
+ from tokenizers import Encoding as EncodingFast
+else:
+
+ @dataclass(frozen=False, eq=True)
+ class AddedToken:
+ """
+ AddedToken represents a token to be added to a Tokenizer An AddedToken can have special options defining the
+ way it should behave.
+
+ The `normalized` will default to `not special` if it is not specified, similarly to the definition in
+ `tokenizers`.
+ """
+
+ def __init__(
+ self, content: str, single_word=False, lstrip=False, rstrip=False, special=False, normalized=None
+ ):
+ self.content = content
+ self.single_word = single_word
+ self.lstrip = lstrip
+ self.rstrip = rstrip
+ self.special = special
+ self.normalized = normalized if normalized is not None else not special
+
+ def __getstate__(self):
+ return self.__dict__
+
+ def __str__(self):
+ return self.content
+
+ @dataclass
+ class EncodingFast:
+ """This is dummy class because without the `tokenizers` library we don't have these objects anyway"""
+
+ pass
+
+
+logger = logging.get_logger(__name__)
+
+VERY_LARGE_INTEGER = int(1e30) # This is used to set the max input length for a model with infinite size input
+LARGE_INTEGER = int(1e20) # This is used when we need something big but slightly smaller than VERY_LARGE_INTEGER
+
+# Define type aliases and NamedTuples
+TextInput = str
+PreTokenizedInput = List[str]
+EncodedInput = List[int]
+TextInputPair = Tuple[str, str]
+PreTokenizedInputPair = Tuple[List[str], List[str]]
+EncodedInputPair = Tuple[List[int], List[int]]
+
+# Define type aliases for text-related non-text modalities
+AudioInput = Union["np.ndarray", "torch.Tensor", List["np.ndarray"], List["torch.Tensor"]]
+
+# Slow tokenizers used to be saved in three separated files
+SPECIAL_TOKENS_MAP_FILE = "special_tokens_map.json"
+ADDED_TOKENS_FILE = "added_tokens.json"
+TOKENIZER_CONFIG_FILE = "tokenizer_config.json"
+CHAT_TEMPLATE_FILE = "chat_template.jinja"
+
+# Fast tokenizers (provided by HuggingFace tokenizer's library) can be saved in a single file
+FULL_TOKENIZER_FILE = "tokenizer.json"
+_re_tokenizer_file = re.compile(r"tokenizer\.(.*)\.json")
+
+
+class TruncationStrategy(ExplicitEnum):
+ """
+ Possible values for the `truncation` argument in [`PreTrainedTokenizerBase.__call__`]. Useful for tab-completion in
+ an IDE.
+ """
+
+ ONLY_FIRST = "only_first"
+ ONLY_SECOND = "only_second"
+ LONGEST_FIRST = "longest_first"
+ DO_NOT_TRUNCATE = "do_not_truncate"
+
+
+class CharSpan(NamedTuple):
+ """
+ Character span in the original string.
+
+ Args:
+ start (`int`): Index of the first character in the original string.
+ end (`int`): Index of the character following the last character in the original string.
+ """
+
+ start: int
+ end: int
+
+
+class TokenSpan(NamedTuple):
+ """
+ Token span in an encoded string (list of tokens).
+
+ Args:
+ start (`int`): Index of the first token in the span.
+ end (`int`): Index of the token following the last token in the span.
+ """
+
+ start: int
+ end: int
+
+
+class BatchEncoding(UserDict):
+ """
+ Holds the output of the [`~tokenization_utils_base.PreTrainedTokenizerBase.__call__`],
+ [`~tokenization_utils_base.PreTrainedTokenizerBase.encode_plus`] and
+ [`~tokenization_utils_base.PreTrainedTokenizerBase.batch_encode_plus`] methods (tokens, attention_masks, etc).
+
+ This class is derived from a python dictionary and can be used as a dictionary. In addition, this class exposes
+ utility methods to map from word/character space to token space.
+
+ Args:
+ data (`dict`, *optional*):
+ Dictionary of lists/arrays/tensors returned by the `__call__`/`encode_plus`/`batch_encode_plus` methods
+ ('input_ids', 'attention_mask', etc.).
+ encoding (`tokenizers.Encoding` or `Sequence[tokenizers.Encoding]`, *optional*):
+ If the tokenizer is a fast tokenizer which outputs additional information like mapping from word/character
+ space to token space the `tokenizers.Encoding` instance or list of instance (for batches) hold this
+ information.
+ tensor_type (`Union[None, str, TensorType]`, *optional*):
+ You can give a tensor_type here to convert the lists of integers in PyTorch/TensorFlow/Numpy Tensors at
+ initialization.
+ prepend_batch_axis (`bool`, *optional*, defaults to `False`):
+ Whether or not to add a batch axis when converting to tensors (see `tensor_type` above). Note that this
+ parameter has an effect if the parameter `tensor_type` is set, *otherwise has no effect*.
+ n_sequences (`Optional[int]`, *optional*):
+ You can give a tensor_type here to convert the lists of integers in PyTorch/TensorFlow/Numpy Tensors at
+ initialization.
+ """
+
+ def __init__(
+ self,
+ data: Optional[Dict[str, Any]] = None,
+ encoding: Optional[Union[EncodingFast, Sequence[EncodingFast]]] = None,
+ tensor_type: Union[None, str, TensorType] = None,
+ prepend_batch_axis: bool = False,
+ n_sequences: Optional[int] = None,
+ ):
+ super().__init__(data)
+
+ if isinstance(encoding, EncodingFast):
+ encoding = [encoding]
+
+ self._encodings = encoding
+
+ if n_sequences is None and encoding is not None and len(encoding):
+ n_sequences = encoding[0].n_sequences
+
+ self._n_sequences = n_sequences
+
+ self.convert_to_tensors(tensor_type=tensor_type, prepend_batch_axis=prepend_batch_axis)
+
+ @property
+ def n_sequences(self) -> Optional[int]:
+ """
+ `Optional[int]`: The number of sequences used to generate each sample from the batch encoded in this
+ [`BatchEncoding`]. Currently can be one of `None` (unknown), `1` (a single sentence) or `2` (a pair of
+ sentences)
+ """
+ return self._n_sequences
+
+ @property
+ def is_fast(self) -> bool:
+ """
+ `bool`: Indicate whether this [`BatchEncoding`] was generated from the result of a [`PreTrainedTokenizerFast`]
+ or not.
+ """
+ return self._encodings is not None
+
+ def __getitem__(self, item: Union[int, str]) -> Union[Any, EncodingFast]:
+ """
+ If the key is a string, returns the value of the dict associated to `key` ('input_ids', 'attention_mask',
+ etc.).
+
+ If the key is an integer, get the `tokenizers.Encoding` for batch item with index `key`.
+
+ If the key is a slice, returns the value of the dict associated to `key` ('input_ids', 'attention_mask', etc.)
+ with the constraint of slice.
+ """
+ if isinstance(item, str):
+ return self.data[item]
+ elif self._encodings is not None:
+ return self._encodings[item]
+ elif isinstance(item, slice):
+ return {key: self.data[key][item] for key in self.data.keys()}
+ else:
+ raise KeyError(
+ "Invalid key. Only three types of key are available: "
+ "(1) string, (2) integers for backend Encoding, and (3) slices for data subsetting."
+ )
+
+ def __getattr__(self, item: str):
+ try:
+ return self.data[item]
+ except KeyError:
+ raise AttributeError
+
+ def __getstate__(self):
+ return {"data": self.data, "encodings": self._encodings}
+
+ def __setstate__(self, state):
+ if "data" in state:
+ self.data = state["data"]
+
+ if "encodings" in state:
+ self._encodings = state["encodings"]
+
+ def keys(self):
+ return self.data.keys()
+
+ def values(self):
+ return self.data.values()
+
+ def items(self):
+ return self.data.items()
+
+ # After this point:
+ # Extended properties and methods only available for fast (Rust-based) tokenizers
+ # provided by HuggingFace tokenizers library.
+
+ @property
+ def encodings(self) -> Optional[List[EncodingFast]]:
+ """
+ `Optional[List[tokenizers.Encoding]]`: The list all encodings from the tokenization process. Returns `None` if
+ the input was tokenized through Python (i.e., not a fast) tokenizer.
+ """
+ return self._encodings
+
+ def tokens(self, batch_index: int = 0) -> List[str]:
+ """
+ Return the list of tokens (sub-parts of the input strings after word/subword splitting and before conversion to
+ integer indices) at a given batch index (only works for the output of a fast tokenizer).
+
+ Args:
+ batch_index (`int`, *optional*, defaults to 0): The index to access in the batch.
+
+ Returns:
+ `List[str]`: The list of tokens at that index.
+ """
+ if not self._encodings:
+ raise ValueError(
+ "tokens() is not available when using non-fast tokenizers (e.g. instance of a `XxxTokenizerFast`"
+ " class)."
+ )
+ return self._encodings[batch_index].tokens
+
+ def sequence_ids(self, batch_index: int = 0) -> List[Optional[int]]:
+ """
+ Return a list mapping the tokens to the id of their original sentences:
+
+ - `None` for special tokens added around or between sequences,
+ - `0` for tokens corresponding to words in the first sequence,
+ - `1` for tokens corresponding to words in the second sequence when a pair of sequences was jointly
+ encoded.
+
+ Args:
+ batch_index (`int`, *optional*, defaults to 0): The index to access in the batch.
+
+ Returns:
+ `List[Optional[int]]`: A list indicating the sequence id corresponding to each token. Special tokens added
+ by the tokenizer are mapped to `None` and other tokens are mapped to the index of their corresponding
+ sequence.
+ """
+ if not self._encodings:
+ raise ValueError(
+ "sequence_ids() is not available when using non-fast tokenizers (e.g. instance of a `XxxTokenizerFast`"
+ " class)."
+ )
+ return self._encodings[batch_index].sequence_ids
+
+ def words(self, batch_index: int = 0) -> List[Optional[int]]:
+ """
+ Return a list mapping the tokens to their actual word in the initial sentence for a fast tokenizer.
+
+ Args:
+ batch_index (`int`, *optional*, defaults to 0): The index to access in the batch.
+
+ Returns:
+ `List[Optional[int]]`: A list indicating the word corresponding to each token. Special tokens added by the
+ tokenizer are mapped to `None` and other tokens are mapped to the index of their corresponding word
+ (several tokens will be mapped to the same word index if they are parts of that word).
+ """
+ if not self._encodings:
+ raise ValueError(
+ "words() is not available when using non-fast tokenizers (e.g. instance of a `XxxTokenizerFast`"
+ " class)."
+ )
+ warnings.warn(
+ "`BatchEncoding.words()` property is deprecated and should be replaced with the identical, "
+ "but more self-explanatory `BatchEncoding.word_ids()` property.",
+ FutureWarning,
+ )
+ return self.word_ids(batch_index)
+
+ def word_ids(self, batch_index: int = 0) -> List[Optional[int]]:
+ """
+ Return a list mapping the tokens to their actual word in the initial sentence for a fast tokenizer.
+
+ Args:
+ batch_index (`int`, *optional*, defaults to 0): The index to access in the batch.
+
+ Returns:
+ `List[Optional[int]]`: A list indicating the word corresponding to each token. Special tokens added by the
+ tokenizer are mapped to `None` and other tokens are mapped to the index of their corresponding word
+ (several tokens will be mapped to the same word index if they are parts of that word).
+ """
+ if not self._encodings:
+ raise ValueError(
+ "word_ids() is not available when using non-fast tokenizers (e.g. instance of a `XxxTokenizerFast`"
+ " class)."
+ )
+ return self._encodings[batch_index].word_ids
+
+ def token_to_sequence(self, batch_or_token_index: int, token_index: Optional[int] = None) -> int:
+ """
+ Get the index of the sequence represented by the given token. In the general use case, this method returns `0`
+ for a single sequence or the first sequence of a pair, and `1` for the second sequence of a pair
+
+ Can be called as:
+
+ - `self.token_to_sequence(token_index)` if batch size is 1
+ - `self.token_to_sequence(batch_index, token_index)` if batch size is greater than 1
+
+ This method is particularly suited when the input sequences are provided as pre-tokenized sequences (i.e.,
+ words are defined by the user). In this case it allows to easily associate encoded tokens with provided
+ tokenized words.
+
+ Args:
+ batch_or_token_index (`int`):
+ Index of the sequence in the batch. If the batch only comprises one sequence, this can be the index of
+ the token in the sequence.
+ token_index (`int`, *optional*):
+ If a batch index is provided in *batch_or_token_index*, this can be the index of the token in the
+ sequence.
+
+ Returns:
+ `int`: Index of the word in the input sequence.
+ """
+
+ if not self._encodings:
+ raise ValueError("token_to_sequence() is not available when using Python based tokenizers")
+ if token_index is not None:
+ batch_index = batch_or_token_index
+ else:
+ batch_index = 0
+ token_index = batch_or_token_index
+ if batch_index < 0:
+ batch_index = self._batch_size + batch_index
+ if token_index < 0:
+ token_index = self._seq_len + token_index
+ return self._encodings[batch_index].token_to_sequence(token_index)
+
+ def token_to_word(self, batch_or_token_index: int, token_index: Optional[int] = None) -> int:
+ """
+ Get the index of the word corresponding (i.e. comprising) to an encoded token in a sequence of the batch.
+
+ Can be called as:
+
+ - `self.token_to_word(token_index)` if batch size is 1
+ - `self.token_to_word(batch_index, token_index)` if batch size is greater than 1
+
+ This method is particularly suited when the input sequences are provided as pre-tokenized sequences (i.e.,
+ words are defined by the user). In this case it allows to easily associate encoded tokens with provided
+ tokenized words.
+
+ Args:
+ batch_or_token_index (`int`):
+ Index of the sequence in the batch. If the batch only comprise one sequence, this can be the index of
+ the token in the sequence.
+ token_index (`int`, *optional*):
+ If a batch index is provided in *batch_or_token_index*, this can be the index of the token in the
+ sequence.
+
+ Returns:
+ `int`: Index of the word in the input sequence.
+ """
+
+ if not self._encodings:
+ raise ValueError("token_to_word() is not available when using Python based tokenizers")
+ if token_index is not None:
+ batch_index = batch_or_token_index
+ else:
+ batch_index = 0
+ token_index = batch_or_token_index
+ if batch_index < 0:
+ batch_index = self._batch_size + batch_index
+ if token_index < 0:
+ token_index = self._seq_len + token_index
+ return self._encodings[batch_index].token_to_word(token_index)
+
+ def word_to_tokens(
+ self, batch_or_word_index: int, word_index: Optional[int] = None, sequence_index: int = 0
+ ) -> Optional[TokenSpan]:
+ """
+ Get the encoded token span corresponding to a word in a sequence of the batch.
+
+ Token spans are returned as a [`~tokenization_utils_base.TokenSpan`] with:
+
+ - **start** -- Index of the first token.
+ - **end** -- Index of the token following the last token.
+
+ Can be called as:
+
+ - `self.word_to_tokens(word_index, sequence_index: int = 0)` if batch size is 1
+ - `self.word_to_tokens(batch_index, word_index, sequence_index: int = 0)` if batch size is greater or equal to
+ 1
+
+ This method is particularly suited when the input sequences are provided as pre-tokenized sequences (i.e. words
+ are defined by the user). In this case it allows to easily associate encoded tokens with provided tokenized
+ words.
+
+ Args:
+ batch_or_word_index (`int`):
+ Index of the sequence in the batch. If the batch only comprises one sequence, this can be the index of
+ the word in the sequence.
+ word_index (`int`, *optional*):
+ If a batch index is provided in *batch_or_token_index*, this can be the index of the word in the
+ sequence.
+ sequence_index (`int`, *optional*, defaults to 0):
+ If pair of sequences are encoded in the batch this can be used to specify which sequence in the pair (0
+ or 1) the provided word index belongs to.
+
+ Returns:
+ ([`~tokenization_utils_base.TokenSpan`], *optional*): Span of tokens in the encoded sequence. Returns
+ `None` if no tokens correspond to the word. This can happen especially when the token is a special token
+ that has been used to format the tokenization. For example when we add a class token at the very beginning
+ of the tokenization.
+ """
+
+ if not self._encodings:
+ raise ValueError("word_to_tokens() is not available when using Python based tokenizers")
+ if word_index is not None:
+ batch_index = batch_or_word_index
+ else:
+ batch_index = 0
+ word_index = batch_or_word_index
+ if batch_index < 0:
+ batch_index = self._batch_size + batch_index
+ if word_index < 0:
+ word_index = self._seq_len + word_index
+ span = self._encodings[batch_index].word_to_tokens(word_index, sequence_index)
+ return TokenSpan(*span) if span is not None else None
+
+ def token_to_chars(self, batch_or_token_index: int, token_index: Optional[int] = None) -> Optional[CharSpan]:
+ """
+ Get the character span corresponding to an encoded token in a sequence of the batch.
+
+ Character spans are returned as a [`~tokenization_utils_base.CharSpan`] with:
+
+ - **start** -- Index of the first character in the original string associated to the token.
+ - **end** -- Index of the character following the last character in the original string associated to the
+ token.
+
+ Can be called as:
+
+ - `self.token_to_chars(token_index)` if batch size is 1
+ - `self.token_to_chars(batch_index, token_index)` if batch size is greater or equal to 1
+
+ Args:
+ batch_or_token_index (`int`):
+ Index of the sequence in the batch. If the batch only comprise one sequence, this can be the index of
+ the token in the sequence.
+ token_index (`int`, *optional*):
+ If a batch index is provided in *batch_or_token_index*, this can be the index of the token or tokens in
+ the sequence.
+
+ Returns:
+ [`~tokenization_utils_base.CharSpan`]: Span of characters in the original string, or None, if the token
+ (e.g. , ) doesn't correspond to any chars in the origin string.
+ """
+
+ if not self._encodings:
+ raise ValueError("token_to_chars() is not available when using Python based tokenizers")
+ if token_index is not None:
+ batch_index = batch_or_token_index
+ else:
+ batch_index = 0
+ token_index = batch_or_token_index
+ span_indices = self._encodings[batch_index].token_to_chars(token_index)
+
+ return CharSpan(*span_indices) if span_indices is not None else None
+
+ def char_to_token(
+ self, batch_or_char_index: int, char_index: Optional[int] = None, sequence_index: int = 0
+ ) -> int:
+ """
+ Get the index of the token in the encoded output comprising a character in the original string for a sequence
+ of the batch.
+
+ Can be called as:
+
+ - `self.char_to_token(char_index)` if batch size is 1
+ - `self.char_to_token(batch_index, char_index)` if batch size is greater or equal to 1
+
+ This method is particularly suited when the input sequences are provided as pre-tokenized sequences (i.e. words
+ are defined by the user). In this case it allows to easily associate encoded tokens with provided tokenized
+ words.
+
+ Args:
+ batch_or_char_index (`int`):
+ Index of the sequence in the batch. If the batch only comprise one sequence, this can be the index of
+ the word in the sequence
+ char_index (`int`, *optional*):
+ If a batch index is provided in *batch_or_token_index*, this can be the index of the word in the
+ sequence.
+ sequence_index (`int`, *optional*, defaults to 0):
+ If pair of sequences are encoded in the batch this can be used to specify which sequence in the pair (0
+ or 1) the provided character index belongs to.
+
+
+ Returns:
+ `int`: Index of the token, or None if the char index refers to a whitespace only token and whitespace is
+ trimmed with `trim_offsets=True`.
+ """
+
+ if not self._encodings:
+ raise ValueError("char_to_token() is not available when using Python based tokenizers")
+ if char_index is not None:
+ batch_index = batch_or_char_index
+ else:
+ batch_index = 0
+ char_index = batch_or_char_index
+ return self._encodings[batch_index].char_to_token(char_index, sequence_index)
+
+ def word_to_chars(
+ self, batch_or_word_index: int, word_index: Optional[int] = None, sequence_index: int = 0
+ ) -> CharSpan:
+ """
+ Get the character span in the original string corresponding to given word in a sequence of the batch.
+
+ Character spans are returned as a CharSpan NamedTuple with:
+
+ - start: index of the first character in the original string
+ - end: index of the character following the last character in the original string
+
+ Can be called as:
+
+ - `self.word_to_chars(word_index)` if batch size is 1
+ - `self.word_to_chars(batch_index, word_index)` if batch size is greater or equal to 1
+
+ Args:
+ batch_or_word_index (`int`):
+ Index of the sequence in the batch. If the batch only comprise one sequence, this can be the index of
+ the word in the sequence
+ word_index (`int`, *optional*):
+ If a batch index is provided in *batch_or_token_index*, this can be the index of the word in the
+ sequence.
+ sequence_index (`int`, *optional*, defaults to 0):
+ If pair of sequences are encoded in the batch this can be used to specify which sequence in the pair (0
+ or 1) the provided word index belongs to.
+
+ Returns:
+ `CharSpan` or `List[CharSpan]`: Span(s) of the associated character or characters in the string. CharSpan
+ are NamedTuple with:
+
+ - start: index of the first character associated to the token in the original string
+ - end: index of the character following the last character associated to the token in the original
+ string
+ """
+
+ if not self._encodings:
+ raise ValueError("word_to_chars() is not available when using Python based tokenizers")
+ if word_index is not None:
+ batch_index = batch_or_word_index
+ else:
+ batch_index = 0
+ word_index = batch_or_word_index
+ return CharSpan(*(self._encodings[batch_index].word_to_chars(word_index, sequence_index)))
+
+ def char_to_word(self, batch_or_char_index: int, char_index: Optional[int] = None, sequence_index: int = 0) -> int:
+ """
+ Get the word in the original string corresponding to a character in the original string of a sequence of the
+ batch.
+
+ Can be called as:
+
+ - `self.char_to_word(char_index)` if batch size is 1
+ - `self.char_to_word(batch_index, char_index)` if batch size is greater than 1
+
+ This method is particularly suited when the input sequences are provided as pre-tokenized sequences (i.e. words
+ are defined by the user). In this case it allows to easily associate encoded tokens with provided tokenized
+ words.
+
+ Args:
+ batch_or_char_index (`int`):
+ Index of the sequence in the batch. If the batch only comprise one sequence, this can be the index of
+ the character in the original string.
+ char_index (`int`, *optional*):
+ If a batch index is provided in *batch_or_token_index*, this can be the index of the character in the
+ original string.
+ sequence_index (`int`, *optional*, defaults to 0):
+ If pair of sequences are encoded in the batch this can be used to specify which sequence in the pair (0
+ or 1) the provided character index belongs to.
+
+
+ Returns:
+ `int` or `List[int]`: Index or indices of the associated encoded token(s).
+ """
+
+ if not self._encodings:
+ raise ValueError("char_to_word() is not available when using Python based tokenizers")
+ if char_index is not None:
+ batch_index = batch_or_char_index
+ else:
+ batch_index = 0
+ char_index = batch_or_char_index
+ return self._encodings[batch_index].char_to_word(char_index, sequence_index)
+
+ def convert_to_tensors(
+ self, tensor_type: Optional[Union[str, TensorType]] = None, prepend_batch_axis: bool = False
+ ):
+ """
+ Convert the inner content to tensors.
+
+ Args:
+ tensor_type (`str` or [`~utils.TensorType`], *optional*):
+ The type of tensors to use. If `str`, should be one of the values of the enum [`~utils.TensorType`]. If
+ `None`, no modification is done.
+ prepend_batch_axis (`int`, *optional*, defaults to `False`):
+ Whether or not to add the batch dimension during the conversion.
+ """
+ if tensor_type is None:
+ return self
+
+ # Convert to TensorType
+ if not isinstance(tensor_type, TensorType):
+ tensor_type = TensorType(tensor_type)
+
+ # Get a function reference for the correct framework
+ if tensor_type == TensorType.TENSORFLOW:
+ if not is_tf_available():
+ raise ImportError(
+ "Unable to convert output to TensorFlow tensors format, TensorFlow is not installed."
+ )
+ import tensorflow as tf
+
+ as_tensor = tf.constant
+ is_tensor = tf.is_tensor
+ elif tensor_type == TensorType.PYTORCH:
+ if not is_torch_available():
+ raise ImportError("Unable to convert output to PyTorch tensors format, PyTorch is not installed.")
+ import torch
+
+ is_tensor = torch.is_tensor
+
+ def as_tensor(value, dtype=None):
+ if isinstance(value, list) and isinstance(value[0], np.ndarray):
+ return torch.from_numpy(np.array(value))
+ return torch.tensor(value)
+
+ elif tensor_type == TensorType.JAX:
+ if not is_flax_available():
+ raise ImportError("Unable to convert output to JAX tensors format, JAX is not installed.")
+ import jax.numpy as jnp # noqa: F811
+
+ as_tensor = jnp.array
+ is_tensor = is_jax_tensor
+
+ elif tensor_type == TensorType.MLX:
+ if not is_mlx_available():
+ raise ImportError("Unable to convert output to MLX tensors format, MLX is not installed.")
+ import mlx.core as mx
+
+ as_tensor = mx.array
+
+ def is_tensor(obj):
+ return isinstance(obj, mx.array)
+ else:
+
+ def as_tensor(value, dtype=None):
+ if isinstance(value, (list, tuple)) and isinstance(value[0], (list, tuple, np.ndarray)):
+ value_lens = [len(val) for val in value]
+ if len(set(value_lens)) > 1 and dtype is None:
+ # we have a ragged list so handle explicitly
+ value = as_tensor([np.asarray(val) for val in value], dtype=object)
+ return np.asarray(value, dtype=dtype)
+
+ is_tensor = is_numpy_array
+
+ # Do the tensor conversion in batch
+ for key, value in self.items():
+ try:
+ if prepend_batch_axis:
+ value = [value]
+
+ if not is_tensor(value):
+ tensor = as_tensor(value)
+
+ # Removing this for now in favor of controlling the shape with `prepend_batch_axis`
+ # # at-least2d
+ # if tensor.ndim > 2:
+ # tensor = tensor.squeeze(0)
+ # elif tensor.ndim < 2:
+ # tensor = tensor[None, :]
+
+ self[key] = tensor
+ except Exception as e:
+ if key == "overflowing_tokens":
+ raise ValueError(
+ "Unable to create tensor returning overflowing tokens of different lengths. "
+ "Please see if a fast version of this tokenizer is available to have this feature available."
+ ) from e
+ raise ValueError(
+ "Unable to create tensor, you should probably activate truncation and/or padding with"
+ " 'padding=True' 'truncation=True' to have batched tensors with the same length. Perhaps your"
+ f" features (`{key}` in this case) have excessive nesting (inputs type `list` where type `int` is"
+ " expected)."
+ ) from e
+
+ return self
+
+ def to(self, device: Union[str, "torch.device"], *, non_blocking: bool = False) -> "BatchEncoding":
+ """
+ Send all values to device by calling `v.to(device, non_blocking=non_blocking)` (PyTorch only).
+
+ Args:
+ device (`str` or `torch.device`): The device to put the tensors on.
+ non_blocking (`bool`): Whether to perform the copy asynchronously.
+
+ Returns:
+ [`BatchEncoding`]: The same instance after modification.
+ """
+ requires_backends(self, ["torch"])
+ import torch
+
+ # This check catches things like APEX blindly calling "to" on all inputs to a module
+ # Otherwise it passes the casts down and casts the LongTensor containing the token idxs
+ # into a HalfTensor
+ if isinstance(device, str) or is_torch_device(device) or isinstance(device, int):
+ self.data = {
+ k: v.to(device=device, non_blocking=non_blocking) if isinstance(v, torch.Tensor) else v
+ for k, v in self.data.items()
+ }
+ else:
+ logger.warning(f"Attempting to cast a BatchEncoding to type {str(device)}. This is not supported.")
+ return self
+
+
+class SpecialTokensMixin:
+ """
+ A mixin derived by [`PreTrainedTokenizer`] and [`PreTrainedTokenizerFast`] to handle specific behaviors related to
+ special tokens. In particular, this class hold the attributes which can be used to directly access these special
+ tokens in a model-independent manner and allow to set and update the special tokens.
+
+ Args:
+ bos_token (`str` or `tokenizers.AddedToken`, *optional*):
+ A special token representing the beginning of a sentence.
+ eos_token (`str` or `tokenizers.AddedToken`, *optional*):
+ A special token representing the end of a sentence.
+ unk_token (`str` or `tokenizers.AddedToken`, *optional*):
+ A special token representing an out-of-vocabulary token.
+ sep_token (`str` or `tokenizers.AddedToken`, *optional*):
+ A special token separating two different sentences in the same input (used by BERT for instance).
+ pad_token (`str` or `tokenizers.AddedToken`, *optional*):
+ A special token used to make arrays of tokens the same size for batching purpose. Will then be ignored by
+ attention mechanisms or loss computation.
+ cls_token (`str` or `tokenizers.AddedToken`, *optional*):
+ A special token representing the class of the input (used by BERT for instance).
+ mask_token (`str` or `tokenizers.AddedToken`, *optional*):
+ A special token representing a masked token (used by masked-language modeling pretraining objectives, like
+ BERT).
+ additional_special_tokens (tuple or list of `str` or `tokenizers.AddedToken`, *optional*):
+ A tuple or a list of additional tokens, which will be marked as `special`, meaning that they will be
+ skipped when decoding if `skip_special_tokens` is set to `True`.
+ """
+
+ SPECIAL_TOKENS_ATTRIBUTES = [
+ "bos_token",
+ "eos_token",
+ "unk_token",
+ "sep_token",
+ "pad_token",
+ "cls_token",
+ "mask_token",
+ "additional_special_tokens",
+ ]
+
+ def __init__(self, verbose=False, **kwargs):
+ self._pad_token_type_id = 0
+ self.verbose = verbose
+ self._special_tokens_map = dict.fromkeys(self.SPECIAL_TOKENS_ATTRIBUTES)
+ self._special_tokens_map["additional_special_tokens"] = [] # for BC where it defaults to empty list
+
+ # We directly set the hidden value to allow initialization with special tokens
+ # which are not yet in the vocabulary. Necessary for serialization/de-serialization
+ # TODO clean this up at some point (probably by switching to fast tokenizers)
+
+ for key, value in kwargs.items():
+ if value is None:
+ continue
+ if key in self.SPECIAL_TOKENS_ATTRIBUTES:
+ if key == "additional_special_tokens":
+ assert isinstance(value, (list, tuple)), f"Value {value} is not a list or tuple"
+ assert all(isinstance(t, (str, AddedToken)) for t in value), (
+ "One of the tokens is not a string or an AddedToken"
+ )
+ setattr(self, key, value)
+ elif isinstance(value, (str, AddedToken)):
+ setattr(self, key, value)
+ else:
+ raise TypeError(f"Special token {key} has to be either str or AddedToken but got: {type(value)}")
+
+ def sanitize_special_tokens(self) -> int:
+ """
+ The `sanitize_special_tokens` is now deprecated kept for backward compatibility and will be removed in
+ transformers v5.
+ """
+ logger.warning_once("The `sanitize_special_tokens` will be removed in transformers v5.")
+ return self.add_tokens(self.all_special_tokens_extended, special_tokens=True)
+
+ def add_special_tokens(
+ self, special_tokens_dict: Dict[str, Union[str, AddedToken]], replace_additional_special_tokens=True
+ ) -> int:
+ """
+ Add a dictionary of special tokens (eos, pad, cls, etc.) to the encoder and link them to class attributes. If
+ special tokens are NOT in the vocabulary, they are added to it (indexed starting from the last index of the
+ current vocabulary).
+
+ When adding new tokens to the vocabulary, you should make sure to also resize the token embedding matrix of the
+ model so that its embedding matrix matches the tokenizer.
+
+ In order to do that, please use the [`~PreTrainedModel.resize_token_embeddings`] method.
+
+ Using `add_special_tokens` will ensure your special tokens can be used in several ways:
+
+ - Special tokens can be skipped when decoding using `skip_special_tokens = True`.
+ - Special tokens are carefully handled by the tokenizer (they are never split), similar to `AddedTokens`.
+ - You can easily refer to special tokens using tokenizer class attributes like `tokenizer.cls_token`. This
+ makes it easy to develop model-agnostic training and fine-tuning scripts.
+
+ When possible, special tokens are already registered for provided pretrained models (for instance
+ [`BertTokenizer`] `cls_token` is already registered to be :obj*'[CLS]'* and XLM's one is also registered to be
+ `'
]({self._wandb.run.get_url()})'
+ )
+
+ modelcard.AUTOGENERATED_TRAINER_COMMENT += f"\n{badge_markdown}"
+
+ def on_train_begin(self, args, state, control, model=None, **kwargs):
+ if self._wandb is None:
+ return
+ hp_search = state.is_hyper_param_search
+ if hp_search:
+ self._wandb.finish()
+ self._initialized = False
+ args.run_name = None
+ if not self._initialized:
+ self.setup(args, state, model, **kwargs)
+
+ def on_train_end(self, args, state, control, model=None, processing_class=None, **kwargs):
+ if self._wandb is None:
+ return
+ if self._log_model.is_enabled and self._initialized and state.is_world_process_zero:
+ from ..trainer import Trainer
+
+ fake_trainer = Trainer(args=args, model=model, processing_class=processing_class, eval_dataset=["fake"])
+ with tempfile.TemporaryDirectory() as temp_dir:
+ fake_trainer.save_model(temp_dir)
+ metadata = (
+ {
+ k: v
+ for k, v in dict(self._wandb.summary).items()
+ if isinstance(v, numbers.Number) and not k.startswith("_")
+ }
+ if not args.load_best_model_at_end
+ else {
+ f"eval/{args.metric_for_best_model}": state.best_metric,
+ "train/total_floss": state.total_flos,
+ "model/num_parameters": self._wandb.config.get("model/num_parameters"),
+ }
+ )
+ metadata["final_model"] = True
+ logger.info("Logging model artifacts. ...")
+ model_name = (
+ f"model-{self._wandb.run.id}"
+ if (args.run_name is None or args.run_name == args.output_dir)
+ else f"model-{self._wandb.run.name}"
+ )
+ # add the model architecture to a separate text file
+ save_model_architecture_to_file(model, temp_dir)
+
+ artifact = self._wandb.Artifact(name=model_name, type="model", metadata=metadata)
+ for f in Path(temp_dir).glob("*"):
+ if f.is_file():
+ with artifact.new_file(f.name, mode="wb") as fa:
+ fa.write(f.read_bytes())
+ self._wandb.run.log_artifact(artifact, aliases=["final_model"])
+
+ def on_log(self, args, state, control, model=None, logs=None, **kwargs):
+ single_value_scalars = [
+ "train_runtime",
+ "train_samples_per_second",
+ "train_steps_per_second",
+ "train_loss",
+ "total_flos",
+ ]
+
+ if self._wandb is None:
+ return
+ if not self._initialized:
+ self.setup(args, state, model)
+ if state.is_world_process_zero:
+ for k, v in logs.items():
+ if k in single_value_scalars:
+ self._wandb.run.summary[k] = v
+ non_scalar_logs = {k: v for k, v in logs.items() if k not in single_value_scalars}
+ non_scalar_logs = rewrite_logs(non_scalar_logs)
+ self._wandb.log({**non_scalar_logs, "train/global_step": state.global_step})
+
+ def on_save(self, args, state, control, **kwargs):
+ if self._log_model == WandbLogModel.CHECKPOINT and self._initialized and state.is_world_process_zero:
+ checkpoint_metadata = {
+ k: v
+ for k, v in dict(self._wandb.summary).items()
+ if isinstance(v, numbers.Number) and not k.startswith("_")
+ }
+ checkpoint_metadata["model/num_parameters"] = self._wandb.config.get("model/num_parameters")
+
+ ckpt_dir = f"checkpoint-{state.global_step}"
+ artifact_path = os.path.join(args.output_dir, ckpt_dir)
+ logger.info(f"Logging checkpoint artifacts in {ckpt_dir}. ...")
+ checkpoint_name = (
+ f"model-{self._wandb.run.id}"
+ if (args.run_name is None or args.run_name == args.output_dir)
+ else f"model-{self._wandb.run.name}"
+ )
+ artifact = self._wandb.Artifact(name=checkpoint_name, type="model", metadata=checkpoint_metadata)
+ artifact.add_dir(artifact_path)
+ self._wandb.log_artifact(
+ artifact, aliases=[f"epoch_{round(state.epoch, 2)}", f"checkpoint_global_step_{state.global_step}"]
+ )
+
+ def on_predict(self, args, state, control, metrics, **kwargs):
+ if self._wandb is None:
+ return
+ if not self._initialized:
+ self.setup(args, state, **kwargs)
+ if state.is_world_process_zero:
+ metrics = rewrite_logs(metrics)
+ self._wandb.log(metrics)
+
+
+class CometCallback(TrainerCallback):
+ """
+ A [`TrainerCallback`] that sends the logs to [Comet ML](https://www.comet.com/site/).
+ """
+
+ def __init__(self):
+ if _is_comet_installed is False or _is_comet_recent_enough is False:
+ raise RuntimeError(
+ f"CometCallback requires comet-ml>={_MIN_COMET_VERSION} to be installed. Run `pip install comet-ml>={_MIN_COMET_VERSION}`."
+ )
+ self._initialized = False
+ self._log_assets = False
+ self._experiment = None
+
+ def setup(self, args, state, model):
+ """
+ Setup the optional Comet integration.
+
+ Environment:
+ - **COMET_MODE** (`str`, *optional*, default to `get_or_create`):
+ Control whether to create and log to a new Comet experiment or append to an existing experiment.
+ It accepts the following values:
+ * `get_or_create`: Decides automatically depending if
+ `COMET_EXPERIMENT_KEY` is set and whether an Experiment
+ with that key already exists or not.
+ * `create`: Always create a new Comet Experiment.
+ * `get`: Always try to append to an Existing Comet Experiment.
+ Requires `COMET_EXPERIMENT_KEY` to be set.
+ * `ONLINE`: **deprecated**, used to create an online
+ Experiment. Use `COMET_START_ONLINE=1` instead.
+ * `OFFLINE`: **deprecated**, used to created an offline
+ Experiment. Use `COMET_START_ONLINE=0` instead.
+ * `DISABLED`: **deprecated**, used to disable Comet logging.
+ Use the `--report_to` flag to control the integrations used
+ for logging result instead.
+ - **COMET_PROJECT_NAME** (`str`, *optional*):
+ Comet project name for experiments.
+ - **COMET_LOG_ASSETS** (`str`, *optional*, defaults to `TRUE`):
+ Whether or not to log training assets (tf event logs, checkpoints, etc), to Comet. Can be `TRUE`, or
+ `FALSE`.
+
+ For a number of configurable items in the environment, see
+ [here](https://www.comet.com/docs/v2/guides/experiment-management/configure-sdk/#explore-comet-configuration-options).
+ """
+ self._initialized = True
+ log_assets = os.getenv("COMET_LOG_ASSETS", "FALSE").upper()
+ if log_assets in {"TRUE", "1"}:
+ self._log_assets = True
+ if state.is_world_process_zero:
+ comet_old_mode = os.getenv("COMET_MODE")
+
+ mode = None
+ online = None
+
+ if comet_old_mode is not None:
+ comet_old_mode = comet_old_mode.lower()
+
+ if comet_old_mode == "online":
+ online = True
+ elif comet_old_mode == "offline":
+ online = False
+ elif comet_old_mode in ("get", "get_or_create", "create"):
+ mode = comet_old_mode
+ elif comet_old_mode:
+ logger.warning("Invalid COMET_MODE env value %r, Comet logging is disabled", comet_old_mode)
+ return
+
+ # For HPO, we always create a new experiment for each trial
+ if state.is_hyper_param_search:
+ if mode is not None:
+ logger.warning(
+ "Hyperparameter Search is enabled, forcing the creation of new experimetns, COMET_MODE value %r is ignored",
+ comet_old_mode,
+ )
+ mode = "create"
+
+ import comet_ml
+
+ # Do not use the default run_name as the experiment name
+ if args.run_name is not None and args.run_name != args.output_dir:
+ experiment_config = comet_ml.ExperimentConfig(name=args.run_name)
+ else:
+ experiment_config = comet_ml.ExperimentConfig()
+
+ self._experiment = comet_ml.start(online=online, mode=mode, experiment_config=experiment_config)
+ self._experiment.__internal_api__set_model_graph__(model, framework="transformers")
+
+ params = {"args": args.to_dict()}
+
+ if hasattr(model, "config") and model.config is not None:
+ model_config = model.config.to_dict()
+ params["config"] = model_config
+ if hasattr(model, "peft_config") and model.peft_config is not None:
+ peft_config = model.peft_config
+ params["peft_config"] = peft_config
+
+ self._experiment.__internal_api__log_parameters__(
+ params, framework="transformers", source="manual", flatten_nested=True
+ )
+
+ if state.is_hyper_param_search:
+ optimization_id = getattr(state, "trial_name", None)
+ optimization_params = getattr(state, "trial_params", None)
+
+ self._experiment.log_optimization(optimization_id=optimization_id, parameters=optimization_params)
+
+ def on_train_begin(self, args, state, control, model=None, **kwargs):
+ if not self._initialized:
+ self.setup(args, state, model)
+
+ def on_log(self, args, state, control, model=None, logs=None, **kwargs):
+ if not self._initialized:
+ self.setup(args, state, model)
+ if state.is_world_process_zero:
+ if self._experiment is not None:
+ rewritten_logs = rewrite_logs(logs)
+ self._experiment.__internal_api__log_metrics__(
+ rewritten_logs, step=state.global_step, epoch=state.epoch, framework="transformers"
+ )
+
+ def on_train_end(self, args, state, control, **kwargs):
+ if self._initialized and state.is_world_process_zero:
+ if self._experiment is not None:
+ if self._log_assets is True:
+ logger.info("Logging checkpoints. This may take time.")
+ self._experiment.log_asset_folder(
+ args.output_dir, recursive=True, log_file_name=True, step=state.global_step
+ )
+
+ # We create one experiment per trial in HPO mode
+ if state.is_hyper_param_search:
+ self._experiment.clean()
+ self._initialized = False
+
+ def on_predict(self, args, state, control, metrics, **kwargs):
+ if not self._initialized:
+ self.setup(args, state, model=None)
+ if state.is_world_process_zero and self._experiment is not None:
+ rewritten_metrics = rewrite_logs(metrics)
+ self._experiment.__internal_api__log_metrics__(
+ rewritten_metrics, step=state.global_step, epoch=state.epoch, framework="transformers"
+ )
+
+
+class AzureMLCallback(TrainerCallback):
+ """
+ A [`TrainerCallback`] that sends the logs to [AzureML](https://pypi.org/project/azureml-sdk/).
+ """
+
+ def __init__(self, azureml_run=None):
+ if not is_azureml_available():
+ raise RuntimeError("AzureMLCallback requires azureml to be installed. Run `pip install azureml-sdk`.")
+ self.azureml_run = azureml_run
+
+ def on_init_end(self, args, state, control, **kwargs):
+ from azureml.core.run import Run
+
+ if self.azureml_run is None and state.is_world_process_zero:
+ self.azureml_run = Run.get_context()
+
+ def on_log(self, args, state, control, logs=None, **kwargs):
+ if self.azureml_run and state.is_world_process_zero:
+ for k, v in logs.items():
+ if isinstance(v, (int, float)):
+ self.azureml_run.log(k, v, description=k)
+
+
+class MLflowCallback(TrainerCallback):
+ """
+ A [`TrainerCallback`] that sends the logs to [MLflow](https://www.mlflow.org/). Can be disabled by setting
+ environment variable `DISABLE_MLFLOW_INTEGRATION = TRUE`.
+ """
+
+ def __init__(self):
+ if not is_mlflow_available():
+ raise RuntimeError("MLflowCallback requires mlflow to be installed. Run `pip install mlflow`.")
+ import mlflow
+
+ self._MAX_PARAM_VAL_LENGTH = mlflow.utils.validation.MAX_PARAM_VAL_LENGTH
+ self._MAX_PARAMS_TAGS_PER_BATCH = mlflow.utils.validation.MAX_PARAMS_TAGS_PER_BATCH
+
+ self._initialized = False
+ self._auto_end_run = False
+ self._log_artifacts = False
+ self._ml_flow = mlflow
+
+ def setup(self, args, state, model):
+ """
+ Setup the optional MLflow integration.
+
+ Environment:
+ - **HF_MLFLOW_LOG_ARTIFACTS** (`str`, *optional*):
+ Whether to use MLflow `.log_artifact()` facility to log artifacts. This only makes sense if logging to a
+ remote server, e.g. s3 or GCS. If set to `True` or *1*, will copy each saved checkpoint on each save in
+ [`TrainingArguments`]'s `output_dir` to the local or remote artifact storage. Using it without a remote
+ storage will just copy the files to your artifact location.
+ - **MLFLOW_TRACKING_URI** (`str`, *optional*):
+ Whether to store runs at a specific path or remote server. Unset by default, which skips setting the
+ tracking URI entirely.
+ - **MLFLOW_EXPERIMENT_NAME** (`str`, *optional*, defaults to `None`):
+ Whether to use an MLflow experiment_name under which to launch the run. Default to `None` which will point
+ to the `Default` experiment in MLflow. Otherwise, it is a case sensitive name of the experiment to be
+ activated. If an experiment with this name does not exist, a new experiment with this name is created.
+ - **MLFLOW_TAGS** (`str`, *optional*):
+ A string dump of a dictionary of key/value pair to be added to the MLflow run as tags. Example:
+ `os.environ['MLFLOW_TAGS']='{"release.candidate": "RC1", "release.version": "2.2.0"}'`.
+ - **MLFLOW_NESTED_RUN** (`str`, *optional*):
+ Whether to use MLflow nested runs. If set to `True` or *1*, will create a nested run inside the current
+ run.
+ - **MLFLOW_RUN_ID** (`str`, *optional*):
+ Allow to reattach to an existing run which can be useful when resuming training from a checkpoint. When
+ `MLFLOW_RUN_ID` environment variable is set, `start_run` attempts to resume a run with the specified run ID
+ and other parameters are ignored.
+ - **MLFLOW_FLATTEN_PARAMS** (`str`, *optional*, defaults to `False`):
+ Whether to flatten the parameters dictionary before logging.
+ - **MLFLOW_MAX_LOG_PARAMS** (`int`, *optional*):
+ Set the maximum number of parameters to log in the run.
+ """
+ self._log_artifacts = os.getenv("HF_MLFLOW_LOG_ARTIFACTS", "FALSE").upper() in ENV_VARS_TRUE_VALUES
+ self._nested_run = os.getenv("MLFLOW_NESTED_RUN", "FALSE").upper() in ENV_VARS_TRUE_VALUES
+ self._tracking_uri = os.getenv("MLFLOW_TRACKING_URI", None)
+ self._experiment_name = os.getenv("MLFLOW_EXPERIMENT_NAME", None)
+ self._flatten_params = os.getenv("MLFLOW_FLATTEN_PARAMS", "FALSE").upper() in ENV_VARS_TRUE_VALUES
+ self._run_id = os.getenv("MLFLOW_RUN_ID", None)
+ self._max_log_params = os.getenv("MLFLOW_MAX_LOG_PARAMS", None)
+
+ # "synchronous" flag is only available with mlflow version >= 2.8.0
+ # https://github.com/mlflow/mlflow/pull/9705
+ # https://github.com/mlflow/mlflow/releases/tag/v2.8.0
+ self._async_log = packaging.version.parse(self._ml_flow.__version__) >= packaging.version.parse("2.8.0")
+
+ logger.debug(
+ f"MLflow experiment_name={self._experiment_name}, run_name={args.run_name}, nested={self._nested_run},"
+ f" tracking_uri={self._tracking_uri}"
+ )
+ if state.is_world_process_zero:
+ if not self._ml_flow.is_tracking_uri_set():
+ if self._tracking_uri:
+ self._ml_flow.set_tracking_uri(self._tracking_uri)
+ logger.debug(f"MLflow tracking URI is set to {self._tracking_uri}")
+ else:
+ logger.debug(
+ "Environment variable `MLFLOW_TRACKING_URI` is not provided and therefore will not be"
+ " explicitly set."
+ )
+ else:
+ logger.debug(f"MLflow tracking URI is set to {self._ml_flow.get_tracking_uri()}")
+
+ if self._ml_flow.active_run() is None or self._nested_run or self._run_id:
+ if self._experiment_name:
+ # Use of set_experiment() ensure that Experiment is created if not exists
+ self._ml_flow.set_experiment(self._experiment_name)
+ self._ml_flow.start_run(run_name=args.run_name, nested=self._nested_run)
+ logger.debug(f"MLflow run started with run_id={self._ml_flow.active_run().info.run_id}")
+ self._auto_end_run = True
+ combined_dict = args.to_dict()
+ if hasattr(model, "config") and model.config is not None:
+ model_config = model.config.to_dict()
+ combined_dict = {**model_config, **combined_dict}
+ combined_dict = flatten_dict(combined_dict) if self._flatten_params else combined_dict
+ # remove params that are too long for MLflow
+ for name, value in list(combined_dict.items()):
+ # internally, all values are converted to str in MLflow
+ if len(str(value)) > self._MAX_PARAM_VAL_LENGTH:
+ logger.warning(
+ f'Trainer is attempting to log a value of "{value}" for key "{name}" as a parameter. MLflow\'s'
+ " log_param() only accepts values no longer than 250 characters so we dropped this attribute."
+ " You can use `MLFLOW_FLATTEN_PARAMS` environment variable to flatten the parameters and"
+ " avoid this message."
+ )
+ del combined_dict[name]
+ # MLflow cannot log more than 100 values in one go, so we have to split it
+ combined_dict_items = list(combined_dict.items())
+ if self._max_log_params and self._max_log_params.isdigit():
+ max_log_params = int(self._max_log_params)
+ if max_log_params < len(combined_dict_items):
+ logger.debug(
+ f"Reducing the number of parameters to log from {len(combined_dict_items)} to {max_log_params}."
+ )
+ combined_dict_items = combined_dict_items[:max_log_params]
+ for i in range(0, len(combined_dict_items), self._MAX_PARAMS_TAGS_PER_BATCH):
+ if self._async_log:
+ self._ml_flow.log_params(
+ dict(combined_dict_items[i : i + self._MAX_PARAMS_TAGS_PER_BATCH]), synchronous=False
+ )
+ else:
+ self._ml_flow.log_params(dict(combined_dict_items[i : i + self._MAX_PARAMS_TAGS_PER_BATCH]))
+ mlflow_tags = os.getenv("MLFLOW_TAGS", None)
+ if mlflow_tags:
+ mlflow_tags = json.loads(mlflow_tags)
+ self._ml_flow.set_tags(mlflow_tags)
+ self._initialized = True
+
+ def on_train_begin(self, args, state, control, model=None, **kwargs):
+ if not self._initialized:
+ self.setup(args, state, model)
+
+ def on_log(self, args, state, control, logs, model=None, **kwargs):
+ if not self._initialized:
+ self.setup(args, state, model)
+ if state.is_world_process_zero:
+ metrics = {}
+ for k, v in logs.items():
+ if isinstance(v, (int, float)):
+ metrics[k] = v
+ elif isinstance(v, torch.Tensor) and v.numel() == 1:
+ metrics[k] = v.item()
+ else:
+ logger.warning(
+ f'Trainer is attempting to log a value of "{v}" of type {type(v)} for key "{k}" as a metric. '
+ "MLflow's log_metric() only accepts float and int types so we dropped this attribute."
+ )
+
+ if self._async_log:
+ self._ml_flow.log_metrics(metrics=metrics, step=state.global_step, synchronous=False)
+ else:
+ self._ml_flow.log_metrics(metrics=metrics, step=state.global_step)
+
+ def on_train_end(self, args, state, control, **kwargs):
+ if self._initialized and state.is_world_process_zero:
+ if self._auto_end_run and self._ml_flow.active_run():
+ self._ml_flow.end_run()
+
+ def on_save(self, args, state, control, **kwargs):
+ if self._initialized and state.is_world_process_zero and self._log_artifacts:
+ ckpt_dir = f"checkpoint-{state.global_step}"
+ artifact_path = os.path.join(args.output_dir, ckpt_dir)
+ logger.info(f"Logging checkpoint artifacts in {ckpt_dir}. This may take time.")
+ self._ml_flow.pyfunc.log_model(
+ ckpt_dir,
+ artifacts={"model_path": artifact_path},
+ python_model=self._ml_flow.pyfunc.PythonModel(),
+ )
+
+ def __del__(self):
+ # if the previous run is not terminated correctly, the fluent API will
+ # not let you start a new run before the previous one is killed
+ if (
+ self._auto_end_run
+ and callable(getattr(self._ml_flow, "active_run", None))
+ and self._ml_flow.active_run() is not None
+ ):
+ self._ml_flow.end_run()
+
+
+class DagsHubCallback(MLflowCallback):
+ """
+ A [`TrainerCallback`] that logs to [DagsHub](https://dagshub.com/). Extends [`MLflowCallback`]
+ """
+
+ def __init__(self):
+ super().__init__()
+ if not is_dagshub_available():
+ raise ImportError("DagsHubCallback requires dagshub to be installed. Run `pip install dagshub`.")
+
+ from dagshub.upload import Repo
+
+ self.Repo = Repo
+
+ def setup(self, *args, **kwargs):
+ """
+ Setup the DagsHub's Logging integration.
+
+ Environment:
+ - **HF_DAGSHUB_LOG_ARTIFACTS** (`str`, *optional*):
+ Whether to save the data and model artifacts for the experiment. Default to `False`.
+ """
+
+ self.log_artifacts = os.getenv("HF_DAGSHUB_LOG_ARTIFACTS", "FALSE").upper() in ENV_VARS_TRUE_VALUES
+ self.name = os.getenv("HF_DAGSHUB_MODEL_NAME") or "main"
+ self.remote = os.getenv("MLFLOW_TRACKING_URI")
+ self.repo = self.Repo(
+ owner=self.remote.split(os.sep)[-2],
+ name=self.remote.split(os.sep)[-1].split(".")[0],
+ branch=os.getenv("BRANCH") or "main",
+ )
+ self.path = Path("artifacts")
+
+ if self.remote is None:
+ raise RuntimeError(
+ "DagsHubCallback requires the `MLFLOW_TRACKING_URI` environment variable to be set. Did you run"
+ " `dagshub.init()`?"
+ )
+
+ super().setup(*args, **kwargs)
+
+ def on_train_end(self, args, state, control, **kwargs):
+ if self.log_artifacts:
+ if getattr(self, "train_dataloader", None):
+ torch.save(self.train_dataloader.dataset, os.path.join(args.output_dir, "dataset.pt"))
+
+ self.repo.directory(str(self.path)).add_dir(args.output_dir)
+
+
+class NeptuneMissingConfiguration(Exception):
+ def __init__(self):
+ super().__init__(
+ """
+ ------ Unsupported ---- We were not able to create new runs. You provided a custom Neptune run to
+ `NeptuneCallback` with the `run` argument. For the integration to work fully, provide your `api_token` and
+ `project` by saving them as environment variables or passing them to the callback.
+ """
+ )
+
+
+class NeptuneCallback(TrainerCallback):
+ """TrainerCallback that sends the logs to [Neptune](https://app.neptune.ai).
+
+ Args:
+ api_token (`str`, *optional*): Neptune API token obtained upon registration.
+ You can leave this argument out if you have saved your token to the `NEPTUNE_API_TOKEN` environment
+ variable (strongly recommended). See full setup instructions in the
+ [docs](https://docs.neptune.ai/setup/installation).
+ project (`str`, *optional*): Name of an existing Neptune project, in the form "workspace-name/project-name".
+ You can find and copy the name in Neptune from the project settings -> Properties. If None (default), the
+ value of the `NEPTUNE_PROJECT` environment variable is used.
+ name (`str`, *optional*): Custom name for the run.
+ base_namespace (`str`, *optional*, defaults to "finetuning"): In the Neptune run, the root namespace
+ that will contain all of the metadata logged by the callback.
+ log_parameters (`bool`, *optional*, defaults to `True`):
+ If True, logs all Trainer arguments and model parameters provided by the Trainer.
+ log_checkpoints (`str`, *optional*): If "same", uploads checkpoints whenever they are saved by the Trainer.
+ If "last", uploads only the most recently saved checkpoint. If "best", uploads the best checkpoint (among
+ the ones saved by the Trainer). If `None`, does not upload checkpoints.
+ run (`Run`, *optional*): Pass a Neptune run object if you want to continue logging to an existing run.
+ Read more about resuming runs in the [docs](https://docs.neptune.ai/logging/to_existing_object).
+ **neptune_run_kwargs (*optional*):
+ Additional keyword arguments to be passed directly to the
+ [`neptune.init_run()`](https://docs.neptune.ai/api/neptune#init_run) function when a new run is created.
+
+ For instructions and examples, see the [Transformers integration
+ guide](https://docs.neptune.ai/integrations/transformers) in the Neptune documentation.
+ """
+
+ integration_version_key = "source_code/integrations/transformers"
+ model_parameters_key = "model_parameters"
+ trial_name_key = "trial"
+ trial_params_key = "trial_params"
+ trainer_parameters_key = "trainer_parameters"
+ flat_metrics = {"train/epoch"}
+
+ def __init__(
+ self,
+ *,
+ api_token: Optional[str] = None,
+ project: Optional[str] = None,
+ name: Optional[str] = None,
+ base_namespace: str = "finetuning",
+ run=None,
+ log_parameters: bool = True,
+ log_checkpoints: Optional[str] = None,
+ **neptune_run_kwargs,
+ ):
+ if not is_neptune_available():
+ raise ValueError(
+ "NeptuneCallback requires the Neptune client library to be installed. "
+ "To install the library, run `pip install neptune`."
+ )
+
+ try:
+ from neptune import Run
+ from neptune.internal.utils import verify_type
+ except ImportError:
+ from neptune.new.internal.utils import verify_type
+ from neptune.new.metadata_containers.run import Run
+
+ verify_type("api_token", api_token, (str, type(None)))
+ verify_type("project", project, (str, type(None)))
+ verify_type("name", name, (str, type(None)))
+ verify_type("base_namespace", base_namespace, str)
+ verify_type("run", run, (Run, type(None)))
+ verify_type("log_parameters", log_parameters, bool)
+ verify_type("log_checkpoints", log_checkpoints, (str, type(None)))
+
+ self._base_namespace_path = base_namespace
+ self._log_parameters = log_parameters
+ self._log_checkpoints = log_checkpoints
+ self._initial_run: Optional[Run] = run
+
+ self._run = None
+ self._is_monitoring_run = False
+ self._run_id = None
+ self._force_reset_monitoring_run = False
+ self._init_run_kwargs = {"api_token": api_token, "project": project, "name": name, **neptune_run_kwargs}
+
+ self._volatile_checkpoints_dir = None
+ self._should_upload_checkpoint = self._log_checkpoints is not None
+ self._recent_checkpoint_path = None
+
+ if self._log_checkpoints in {"last", "best"}:
+ self._target_checkpoints_namespace = f"checkpoints/{self._log_checkpoints}"
+ self._should_clean_recently_uploaded_checkpoint = True
+ else:
+ self._target_checkpoints_namespace = "checkpoints"
+ self._should_clean_recently_uploaded_checkpoint = False
+
+ def _stop_run_if_exists(self):
+ if self._run:
+ self._run.stop()
+ del self._run
+ self._run = None
+
+ def _initialize_run(self, **additional_neptune_kwargs):
+ try:
+ from neptune import init_run
+ from neptune.exceptions import NeptuneMissingApiTokenException, NeptuneMissingProjectNameException
+ except ImportError:
+ from neptune.new import init_run
+ from neptune.new.exceptions import NeptuneMissingApiTokenException, NeptuneMissingProjectNameException
+
+ self._stop_run_if_exists()
+
+ try:
+ run_params = additional_neptune_kwargs.copy()
+ run_params.update(self._init_run_kwargs)
+ self._run = init_run(**run_params)
+ self._run_id = self._run["sys/id"].fetch()
+ except (NeptuneMissingProjectNameException, NeptuneMissingApiTokenException) as e:
+ raise NeptuneMissingConfiguration() from e
+
+ def _use_initial_run(self):
+ self._run = self._initial_run
+ self._is_monitoring_run = True
+ self._run_id = self._run["sys/id"].fetch()
+ self._initial_run = None
+
+ def _ensure_run_with_monitoring(self):
+ if self._initial_run is not None:
+ self._use_initial_run()
+ else:
+ if not self._force_reset_monitoring_run and self._is_monitoring_run:
+ return
+
+ if self._run and not self._is_monitoring_run and not self._force_reset_monitoring_run:
+ self._initialize_run(with_id=self._run_id)
+ self._is_monitoring_run = True
+ else:
+ self._initialize_run()
+ self._force_reset_monitoring_run = False
+
+ def _ensure_at_least_run_without_monitoring(self):
+ if self._initial_run is not None:
+ self._use_initial_run()
+ else:
+ if not self._run:
+ self._initialize_run(
+ with_id=self._run_id,
+ capture_stdout=False,
+ capture_stderr=False,
+ capture_hardware_metrics=False,
+ capture_traceback=False,
+ )
+ self._is_monitoring_run = False
+
+ @property
+ def run(self):
+ if self._run is None:
+ self._ensure_at_least_run_without_monitoring()
+ return self._run
+
+ @property
+ def _metadata_namespace(self):
+ return self.run[self._base_namespace_path]
+
+ def _log_integration_version(self):
+ self.run[NeptuneCallback.integration_version_key] = version
+
+ def _log_trainer_parameters(self, args):
+ self._metadata_namespace[NeptuneCallback.trainer_parameters_key] = args.to_sanitized_dict()
+
+ def _log_model_parameters(self, model):
+ from neptune.utils import stringify_unsupported
+
+ if model and hasattr(model, "config") and model.config is not None:
+ self._metadata_namespace[NeptuneCallback.model_parameters_key] = stringify_unsupported(
+ model.config.to_dict()
+ )
+
+ def _log_hyper_param_search_parameters(self, state):
+ if state and hasattr(state, "trial_name"):
+ self._metadata_namespace[NeptuneCallback.trial_name_key] = state.trial_name
+
+ if state and hasattr(state, "trial_params") and state.trial_params is not None:
+ self._metadata_namespace[NeptuneCallback.trial_params_key] = state.trial_params
+
+ def _log_model_checkpoint(self, source_directory: str, checkpoint: str):
+ target_path = relative_path = os.path.join(source_directory, checkpoint)
+
+ if self._volatile_checkpoints_dir is not None:
+ consistent_checkpoint_path = os.path.join(self._volatile_checkpoints_dir, checkpoint)
+ try:
+ # Remove leading ../ from a relative path.
+ cpkt_path = relative_path.replace("..", "").lstrip(os.path.sep)
+ copy_path = os.path.join(consistent_checkpoint_path, cpkt_path)
+ shutil.copytree(relative_path, copy_path)
+ target_path = consistent_checkpoint_path
+ except IOError as e:
+ logger.warning(
+ "NeptuneCallback was unable to made a copy of checkpoint due to I/O exception: '{}'. "
+ "Could fail trying to upload.".format(e)
+ )
+
+ self._metadata_namespace[self._target_checkpoints_namespace].upload_files(target_path)
+
+ if self._should_clean_recently_uploaded_checkpoint and self._recent_checkpoint_path is not None:
+ self._metadata_namespace[self._target_checkpoints_namespace].delete_files(self._recent_checkpoint_path)
+
+ self._recent_checkpoint_path = relative_path
+
+ def on_init_end(self, args, state, control, **kwargs):
+ self._volatile_checkpoints_dir = None
+ if self._log_checkpoints and (args.overwrite_output_dir or args.save_total_limit is not None):
+ self._volatile_checkpoints_dir = tempfile.TemporaryDirectory().name
+
+ if self._log_checkpoints == "best" and not args.load_best_model_at_end:
+ raise ValueError("To save the best model checkpoint, the load_best_model_at_end argument must be enabled.")
+
+ def on_train_begin(self, args, state, control, model=None, **kwargs):
+ if not state.is_world_process_zero:
+ return
+
+ self._ensure_run_with_monitoring()
+ self._force_reset_monitoring_run = True
+
+ self._log_integration_version()
+ if self._log_parameters:
+ self._log_trainer_parameters(args)
+ self._log_model_parameters(model)
+
+ if state.is_hyper_param_search:
+ self._log_hyper_param_search_parameters(state)
+
+ def on_train_end(self, args, state, control, **kwargs):
+ self._stop_run_if_exists()
+
+ def __del__(self):
+ if self._volatile_checkpoints_dir is not None:
+ shutil.rmtree(self._volatile_checkpoints_dir, ignore_errors=True)
+
+ self._stop_run_if_exists()
+
+ def on_save(self, args, state, control, **kwargs):
+ if self._should_upload_checkpoint:
+ self._log_model_checkpoint(args.output_dir, f"checkpoint-{state.global_step}")
+
+ def on_evaluate(self, args, state, control, metrics=None, **kwargs):
+ if self._log_checkpoints == "best":
+ best_metric_name = args.metric_for_best_model
+ if not best_metric_name.startswith("eval_"):
+ best_metric_name = f"eval_{best_metric_name}"
+
+ metric_value = metrics.get(best_metric_name)
+
+ operator = np.greater if args.greater_is_better else np.less
+
+ self._should_upload_checkpoint = state.best_metric is None or operator(metric_value, state.best_metric)
+
+ @classmethod
+ def get_run(cls, trainer):
+ for callback in trainer.callback_handler.callbacks:
+ if isinstance(callback, cls):
+ return callback.run
+
+ raise Exception("The trainer doesn't have a NeptuneCallback configured.")
+
+ def on_log(self, args, state, control, logs: Optional[Dict[str, float]] = None, **kwargs):
+ if not state.is_world_process_zero:
+ return
+
+ if logs is not None:
+ for name, value in rewrite_logs(logs).items():
+ if isinstance(value, (int, float)):
+ if name in NeptuneCallback.flat_metrics:
+ self._metadata_namespace[name] = value
+ else:
+ self._metadata_namespace[name].log(value, step=state.global_step)
+
+
+class CodeCarbonCallback(TrainerCallback):
+ """
+ A [`TrainerCallback`] that tracks the CO2 emission of training.
+ """
+
+ def __init__(self):
+ if not is_codecarbon_available():
+ raise RuntimeError(
+ "CodeCarbonCallback requires `codecarbon` to be installed. Run `pip install codecarbon`."
+ )
+ elif torch.version.hip:
+ raise RuntimeError(
+ "CodeCarbonCallback requires `codecarbon` package, which is not compatible with AMD ROCm (https://github.com/mlco2/codecarbon/pull/490). When using the Trainer, please specify the `report_to` argument (https://huggingface.co/docs/transformers/v4.39.3/en/main_classes/trainer#transformers.TrainingArguments.report_to) to disable CodeCarbonCallback."
+ )
+
+ import codecarbon
+
+ self._codecarbon = codecarbon
+ self.tracker = None
+
+ def on_init_end(self, args, state, control, **kwargs):
+ if self.tracker is None and state.is_local_process_zero:
+ # CodeCarbon will automatically handle environment variables for configuration
+ self.tracker = self._codecarbon.EmissionsTracker(output_dir=args.output_dir)
+
+ def on_train_begin(self, args, state, control, model=None, **kwargs):
+ if self.tracker and state.is_local_process_zero:
+ self.tracker.start()
+
+ def on_train_end(self, args, state, control, **kwargs):
+ if self.tracker and state.is_local_process_zero:
+ self.tracker.stop()
+
+
+class ClearMLCallback(TrainerCallback):
+ """
+ A [`TrainerCallback`] that sends the logs to [ClearML](https://clear.ml/).
+
+ Environment:
+ - **CLEARML_PROJECT** (`str`, *optional*, defaults to `HuggingFace Transformers`):
+ ClearML project name.
+ - **CLEARML_TASK** (`str`, *optional*, defaults to `Trainer`):
+ ClearML task name.
+ - **CLEARML_LOG_MODEL** (`bool`, *optional*, defaults to `False`):
+ Whether to log models as artifacts during training.
+ """
+
+ log_suffix = ""
+
+ _hparams_section = "Transformers"
+ _model_config_section = "Model Configuration"
+ _ignore_hparams_overrides = "_ignore_hparams_ui_overrides_"
+ _ignoge_model_config_overrides = "_ignore_model_config_ui_overrides_"
+ _model_config_description = "The configuration of model number {}."
+ _model_config_description_note = (
+ "Note that, when cloning this task and running it remotely,"
+ " the configuration might be applied to another model instead of this one."
+ " To avoid this, initialize the task externally by calling `Task.init`"
+ " before the `ClearMLCallback` is instantiated."
+ )
+ _train_run_counter = 0
+ _model_connect_counter = 0
+ _task_created_in_callback = False
+ _should_close_on_train_end = None
+
+ def __init__(self):
+ if is_clearml_available():
+ import clearml
+
+ self._clearml = clearml
+ else:
+ raise RuntimeError("ClearMLCallback requires 'clearml' to be installed. Run `pip install clearml`.")
+
+ self._initialized = False
+ self._clearml_task = None
+
+ self._log_model = False
+ self._checkpoints_saved = []
+
+ def setup(self, args, state, model, processing_class, **kwargs):
+ if self._clearml is None:
+ return
+ if self._initialized:
+ return
+ ClearMLCallback._train_run_counter += 1
+ ClearMLCallback._model_connect_counter += 1
+ ClearMLCallback.log_suffix = (
+ "" if ClearMLCallback._train_run_counter == 1 else "_" + str(ClearMLCallback._train_run_counter)
+ )
+ if state.is_world_process_zero:
+ logger.info("Automatic ClearML logging enabled.")
+ if self._clearml_task is None:
+ if ClearMLCallback._should_close_on_train_end is None:
+ if not self._clearml.Task.running_locally() or self._clearml.Task.current_task():
+ ClearMLCallback._should_close_on_train_end = False
+ else:
+ ClearMLCallback._should_close_on_train_end = True
+
+ # This might happen when running inside of a pipeline, where the task is already initialized
+ # from outside of Hugging Face
+ if self._clearml.Task.running_locally() and self._clearml.Task.current_task():
+ self._clearml_task = self._clearml.Task.current_task()
+ self._log_model = os.getenv(
+ "CLEARML_LOG_MODEL",
+ "FALSE" if not ClearMLCallback._task_created_in_callback else "TRUE",
+ ).upper() in ENV_VARS_TRUE_VALUES.union({"TRUE"})
+ logger.info("External ClearML Task has been connected.")
+ else:
+ self._clearml_task = self._clearml.Task.init(
+ project_name=os.getenv("CLEARML_PROJECT", "HuggingFace Transformers"),
+ task_name=os.getenv("CLEARML_TASK", "Trainer"),
+ auto_connect_frameworks={"tensorboard": False, "pytorch": False},
+ output_uri=True,
+ )
+ self._log_model = os.getenv("CLEARML_LOG_MODEL", "TRUE").upper() in ENV_VARS_TRUE_VALUES.union(
+ {"TRUE"}
+ )
+ ClearMLCallback._task_created_in_callback = True
+ logger.info("ClearML Task has been initialized.")
+ self._initialized = True
+
+ suffixed_hparams_section = ClearMLCallback._hparams_section + ClearMLCallback.log_suffix
+ ignore_hparams_config_section = suffixed_hparams_section + "/" + ClearMLCallback._ignore_hparams_overrides
+ if self._clearml.Task.running_locally():
+ self._copy_training_args_as_hparams(args, suffixed_hparams_section)
+ self._clearml_task.set_parameter(
+ name=ignore_hparams_config_section,
+ value=True,
+ value_type=bool,
+ description=(
+ "If True, ignore Transformers hyperparameters overrides done in the UI/backend "
+ + "when running remotely. Otherwise, the overrides will be applied when running remotely"
+ ),
+ )
+ elif not self._clearml_task.get_parameter(ignore_hparams_config_section, default=True, cast=True):
+ self._clearml_task.connect(args, suffixed_hparams_section)
+ else:
+ self._copy_training_args_as_hparams(
+ args, ClearMLCallback._hparams_section + ClearMLCallback.log_suffix
+ )
+
+ if getattr(model, "config", None) is not None:
+ ignore_model_config_section = (
+ suffixed_hparams_section + "/" + ClearMLCallback._ignoge_model_config_overrides
+ )
+ configuration_object_description = ClearMLCallback._model_config_description.format(
+ ClearMLCallback._model_connect_counter
+ )
+ if ClearMLCallback._model_connect_counter != ClearMLCallback._train_run_counter:
+ configuration_object_description += " " + ClearMLCallback._model_config_description_note
+ if self._clearml.Task.running_locally():
+ self._clearml_task.set_parameter(
+ name=ignore_model_config_section,
+ value=True,
+ value_type=bool,
+ description=(
+ "If True, ignore Transformers model configuration overrides done in the UI/backend "
+ + "when running remotely. Otherwise, the overrides will be applied when running remotely"
+ ),
+ )
+ self._clearml_task.set_configuration_object(
+ name=ClearMLCallback._model_config_section + ClearMLCallback.log_suffix,
+ config_dict=model.config.to_dict(),
+ description=configuration_object_description,
+ )
+ elif not self._clearml_task.get_parameter(ignore_model_config_section, default=True, cast=True):
+ model.config = model.config.from_dict(
+ self._clearml_task.get_configuration_object_as_dict(
+ ClearMLCallback._model_config_section + ClearMLCallback.log_suffix
+ )
+ )
+ else:
+ self._clearml_task.set_configuration_object(
+ name=ClearMLCallback._model_config_section + ClearMLCallback.log_suffix,
+ config_dict=model.config.to_dict(),
+ description=configuration_object_description,
+ )
+
+ def on_train_begin(self, args, state, control, model=None, processing_class=None, **kwargs):
+ if self._clearml is None:
+ return
+ self._checkpoints_saved = []
+ if state.is_hyper_param_search:
+ self._initialized = False
+ if not self._initialized:
+ self.setup(args, state, model, processing_class, **kwargs)
+
+ def on_train_end(self, args, state, control, **kwargs):
+ if ClearMLCallback._should_close_on_train_end:
+ self._clearml_task.close()
+ ClearMLCallback._train_run_counter = 0
+
+ def on_log(self, args, state, control, model=None, processing_class=None, logs=None, **kwargs):
+ if self._clearml is None:
+ return
+ if not self._initialized:
+ self.setup(args, state, model, processing_class, **kwargs)
+ if state.is_world_process_zero:
+ eval_prefix = "eval_"
+ eval_prefix_len = len(eval_prefix)
+ test_prefix = "test_"
+ test_prefix_len = len(test_prefix)
+ single_value_scalars = [
+ "train_runtime",
+ "train_samples_per_second",
+ "train_steps_per_second",
+ "train_loss",
+ "total_flos",
+ "epoch",
+ ]
+ for k, v in logs.items():
+ if isinstance(v, (int, float)):
+ if k in single_value_scalars:
+ self._clearml_task.get_logger().report_single_value(
+ name=k + ClearMLCallback.log_suffix, value=v
+ )
+ elif k.startswith(eval_prefix):
+ self._clearml_task.get_logger().report_scalar(
+ title="eval" + ClearMLCallback.log_suffix,
+ series=k[eval_prefix_len:],
+ value=v,
+ iteration=state.global_step,
+ )
+ elif k.startswith(test_prefix):
+ self._clearml_task.get_logger().report_scalar(
+ title="test" + ClearMLCallback.log_suffix,
+ series=k[test_prefix_len:],
+ value=v,
+ iteration=state.global_step,
+ )
+ else:
+ self._clearml_task.get_logger().report_scalar(
+ title="train" + ClearMLCallback.log_suffix,
+ series=k,
+ value=v,
+ iteration=state.global_step,
+ )
+ else:
+ logger.warning(
+ "Trainer is attempting to log a value of "
+ f'"{v}" of type {type(v)} for key "{k}" as a scalar. '
+ "This invocation of ClearML logger's report_scalar() "
+ "is incorrect so we dropped this attribute."
+ )
+
+ def on_save(self, args, state, control, **kwargs):
+ if self._log_model and self._clearml_task and state.is_world_process_zero:
+ ckpt_dir = f"checkpoint-{state.global_step}"
+ artifact_path = os.path.join(args.output_dir, ckpt_dir)
+ name = ckpt_dir + ClearMLCallback.log_suffix
+ logger.info(f"Logging checkpoint artifact `{name}`. This may take some time.")
+ output_model = self._clearml.OutputModel(task=self._clearml_task, name=name)
+ output_model.connect(task=self._clearml_task, name=name)
+ output_model.update_weights_package(
+ weights_path=artifact_path,
+ target_filename=ckpt_dir,
+ iteration=state.global_step,
+ auto_delete_file=False,
+ )
+ self._checkpoints_saved.append(output_model)
+ while args.save_total_limit and args.save_total_limit < len(self._checkpoints_saved):
+ try:
+ self._clearml.model.Model.remove(
+ self._checkpoints_saved[0],
+ delete_weights_file=True,
+ force=True,
+ raise_on_errors=True,
+ )
+ except Exception as e:
+ logger.warning(
+ "Could not remove checkpoint `{}` after going over the `save_total_limit`. Error is: {}".format(
+ self._checkpoints_saved[0].name, e
+ )
+ )
+ break
+ self._checkpoints_saved = self._checkpoints_saved[1:]
+
+ def _copy_training_args_as_hparams(self, training_args, prefix):
+ as_dict = {
+ field.name: getattr(training_args, field.name)
+ for field in fields(training_args)
+ if field.init and not field.name.endswith("_token")
+ }
+ flat_dict = {str(k): v for k, v in self._clearml.utilities.proxy_object.flatten_dictionary(as_dict).items()}
+ self._clearml_task._arguments.copy_from_dict(flat_dict, prefix=prefix)
+
+
+class FlyteCallback(TrainerCallback):
+ """A [`TrainerCallback`] that sends the logs to [Flyte](https://flyte.org/).
+ NOTE: This callback only works within a Flyte task.
+
+ Args:
+ save_log_history (`bool`, *optional*, defaults to `True`):
+ When set to True, the training logs are saved as a Flyte Deck.
+
+ sync_checkpoints (`bool`, *optional*, defaults to `True`):
+ When set to True, checkpoints are synced with Flyte and can be used to resume training in the case of an
+ interruption.
+
+ Example:
+
+ ```python
+ # Note: This example skips over some setup steps for brevity.
+ from flytekit import current_context, task
+
+
+ @task
+ def train_hf_transformer():
+ cp = current_context().checkpoint
+ trainer = Trainer(..., callbacks=[FlyteCallback()])
+ output = trainer.train(resume_from_checkpoint=cp.restore())
+ ```
+ """
+
+ def __init__(self, save_log_history: bool = True, sync_checkpoints: bool = True):
+ super().__init__()
+ if not is_flytekit_available():
+ raise ImportError("FlyteCallback requires flytekit to be installed. Run `pip install flytekit`.")
+
+ if not is_flyte_deck_standard_available() or not is_pandas_available():
+ logger.warning(
+ "Syncing log history requires both flytekitplugins-deck-standard and pandas to be installed. "
+ "Run `pip install flytekitplugins-deck-standard pandas` to enable this feature."
+ )
+ save_log_history = False
+
+ from flytekit import current_context
+
+ self.cp = current_context().checkpoint
+ self.save_log_history = save_log_history
+ self.sync_checkpoints = sync_checkpoints
+
+ def on_save(self, args, state, control, **kwargs):
+ if self.sync_checkpoints and state.is_world_process_zero:
+ ckpt_dir = f"checkpoint-{state.global_step}"
+ artifact_path = os.path.join(args.output_dir, ckpt_dir)
+
+ logger.info(f"Syncing checkpoint in {ckpt_dir} to Flyte. This may take time.")
+ self.cp.save(artifact_path)
+
+ def on_train_end(self, args, state, control, **kwargs):
+ if self.save_log_history:
+ import pandas as pd
+ from flytekit import Deck
+ from flytekitplugins.deck.renderer import TableRenderer
+
+ log_history_df = pd.DataFrame(state.log_history)
+ Deck("Log History", TableRenderer().to_html(log_history_df))
+
+
+class DVCLiveCallback(TrainerCallback):
+ """
+ A [`TrainerCallback`] that sends the logs to [DVCLive](https://www.dvc.org/doc/dvclive).
+
+ Use the environment variables below in `setup` to configure the integration. To customize this callback beyond
+ those environment variables, see [here](https://dvc.org/doc/dvclive/ml-frameworks/huggingface).
+
+ Args:
+ live (`dvclive.Live`, *optional*, defaults to `None`):
+ Optional Live instance. If None, a new instance will be created using **kwargs.
+ log_model (Union[Literal["all"], bool], *optional*, defaults to `None`):
+ Whether to use `dvclive.Live.log_artifact()` to log checkpoints created by [`Trainer`]. If set to `True`,
+ the final checkpoint is logged at the end of training. If set to `"all"`, the entire
+ [`TrainingArguments`]'s `output_dir` is logged at each checkpoint.
+ """
+
+ def __init__(
+ self,
+ live: Optional[Any] = None,
+ log_model: Optional[Union[Literal["all"], bool]] = None,
+ **kwargs,
+ ):
+ if not is_dvclive_available():
+ raise RuntimeError("DVCLiveCallback requires dvclive to be installed. Run `pip install dvclive`.")
+ from dvclive import Live
+
+ self._initialized = False
+ self.live = None
+ if isinstance(live, Live):
+ self.live = live
+ elif live is not None:
+ raise RuntimeError(f"Found class {live.__class__} for live, expected dvclive.Live")
+
+ self._log_model = log_model
+ if self._log_model is None:
+ log_model_env = os.getenv("HF_DVCLIVE_LOG_MODEL", "FALSE")
+ if log_model_env.upper() in ENV_VARS_TRUE_VALUES:
+ self._log_model = True
+ elif log_model_env.lower() == "all":
+ self._log_model = "all"
+
+ def setup(self, args, state, model):
+ """
+ Setup the optional DVCLive integration. To customize this callback beyond the environment variables below, see
+ [here](https://dvc.org/doc/dvclive/ml-frameworks/huggingface).
+
+ Environment:
+ - **HF_DVCLIVE_LOG_MODEL** (`str`, *optional*):
+ Whether to use `dvclive.Live.log_artifact()` to log checkpoints created by [`Trainer`]. If set to `True` or
+ *1*, the final checkpoint is logged at the end of training. If set to `all`, the entire
+ [`TrainingArguments`]'s `output_dir` is logged at each checkpoint.
+ """
+ from dvclive import Live
+
+ self._initialized = True
+ if state.is_world_process_zero:
+ if not self.live:
+ self.live = Live()
+ self.live.log_params(args.to_dict())
+
+ def on_train_begin(self, args, state, control, model=None, **kwargs):
+ if not self._initialized:
+ self.setup(args, state, model)
+
+ def on_log(self, args, state, control, model=None, logs=None, **kwargs):
+ if not self._initialized:
+ self.setup(args, state, model)
+ if state.is_world_process_zero:
+ from dvclive.plots import Metric
+ from dvclive.utils import standardize_metric_name
+
+ for key, value in logs.items():
+ if Metric.could_log(value):
+ self.live.log_metric(standardize_metric_name(key, "dvclive.huggingface"), value)
+ else:
+ logger.warning(
+ "Trainer is attempting to log a value of "
+ f'"{value}" of type {type(value)} for key "{key}" as a scalar. '
+ "This invocation of DVCLive's Live.log_metric() "
+ "is incorrect so we dropped this attribute."
+ )
+ self.live.next_step()
+
+ def on_save(self, args, state, control, **kwargs):
+ if self._log_model == "all" and self._initialized and state.is_world_process_zero:
+ self.live.log_artifact(args.output_dir)
+
+ def on_train_end(self, args, state, control, **kwargs):
+ if self._initialized and state.is_world_process_zero:
+ from transformers.trainer import Trainer
+
+ if self._log_model is True:
+ fake_trainer = Trainer(
+ args=args,
+ model=kwargs.get("model"),
+ processing_class=kwargs.get("processing_class"),
+ eval_dataset=["fake"],
+ )
+ name = "best" if args.load_best_model_at_end else "last"
+ output_dir = os.path.join(args.output_dir, name)
+ fake_trainer.save_model(output_dir)
+ self.live.log_artifact(output_dir, name=name, type="model", copy=True)
+ self.live.end()
+
+
+class SwanLabCallback(TrainerCallback):
+ """
+ A [`TrainerCallback`] that logs metrics, media, model checkpoints to [SwanLab](https://swanlab.cn/).
+ """
+
+ def __init__(self):
+ if not is_swanlab_available():
+ raise RuntimeError("SwanLabCallback requires swanlab to be installed. Run `pip install swanlab`.")
+ import swanlab
+
+ self._swanlab = swanlab
+ self._initialized = False
+ self._log_model = os.getenv("SWANLAB_LOG_MODEL", None)
+
+ def setup(self, args, state, model, **kwargs):
+ """
+ Setup the optional SwanLab (*swanlab*) integration.
+
+ One can subclass and override this method to customize the setup if needed. Find more information
+ [here](https://docs.swanlab.cn/guide_cloud/integration/integration-huggingface-transformers.html).
+
+ You can also override the following environment variables. Find more information about environment
+ variables [here](https://docs.swanlab.cn/en/api/environment-variable.html#environment-variables)
+
+ Environment:
+ - **SWANLAB_API_KEY** (`str`, *optional*, defaults to `None`):
+ Cloud API Key. During login, this environment variable is checked first. If it doesn't exist, the system
+ checks if the user is already logged in. If not, the login process is initiated.
+
+ - If a string is passed to the login interface, this environment variable is ignored.
+ - If the user is already logged in, this environment variable takes precedence over locally stored
+ login information.
+
+ - **SWANLAB_PROJECT** (`str`, *optional*, defaults to `None`):
+ Set this to a custom string to store results in a different project. If not specified, the name of the current
+ running directory is used.
+
+ - **SWANLAB_LOG_DIR** (`str`, *optional*, defaults to `swanlog`):
+ This environment variable specifies the storage path for log files when running in local mode.
+ By default, logs are saved in a folder named swanlog under the working directory.
+
+ - **SWANLAB_MODE** (`Literal["local", "cloud", "disabled"]`, *optional*, defaults to `cloud`):
+ SwanLab's parsing mode, which involves callbacks registered by the operator. Currently, there are three modes:
+ local, cloud, and disabled. Note: Case-sensitive. Find more information
+ [here](https://docs.swanlab.cn/en/api/py-init.html#swanlab-init)
+
+ - **SWANLAB_LOG_MODEL** (`str`, *optional*, defaults to `None`):
+ SwanLab does not currently support the save mode functionality.This feature will be available in a future
+ release
+
+ - **SWANLAB_WEB_HOST** (`str`, *optional*, defaults to `None`):
+ Web address for the SwanLab cloud environment for private version (its free)
+
+ - **SWANLAB_API_HOST** (`str`, *optional*, defaults to `None`):
+ API address for the SwanLab cloud environment for private version (its free)
+
+ """
+ self._initialized = True
+
+ if state.is_world_process_zero:
+ logger.info('Automatic SwanLab logging enabled, to disable set os.environ["SWANLAB_MODE"] = "disabled"')
+ combined_dict = {**args.to_dict()}
+
+ if hasattr(model, "config") and model.config is not None:
+ model_config = model.config if isinstance(model.config, dict) else model.config.to_dict()
+ combined_dict = {**model_config, **combined_dict}
+ if hasattr(model, "peft_config") and model.peft_config is not None:
+ peft_config = model.peft_config
+ combined_dict = {**{"peft_config": peft_config}, **combined_dict}
+ trial_name = state.trial_name
+ init_args = {}
+ if trial_name is not None:
+ init_args["experiment_name"] = f"{args.run_name}-{trial_name}"
+ elif args.run_name is not None:
+ init_args["experiment_name"] = args.run_name
+ init_args["project"] = os.getenv("SWANLAB_PROJECT", None)
+
+ if self._swanlab.get_run() is None:
+ self._swanlab.init(
+ **init_args,
+ )
+ # show transformers logo!
+ self._swanlab.config["FRAMEWORK"] = "🤗transformers"
+ # add config parameters (run may have been created manually)
+ self._swanlab.config.update(combined_dict)
+
+ # add number of model parameters to swanlab config
+ try:
+ self._swanlab.config.update({"model_num_parameters": model.num_parameters()})
+ # get peft model parameters
+ if type(model).__name__ == "PeftModel" or type(model).__name__ == "PeftMixedModel":
+ trainable_params, all_param = model.get_nb_trainable_parameters()
+ self._swanlab.config.update({"peft_model_trainable_params": trainable_params})
+ self._swanlab.config.update({"peft_model_all_param": all_param})
+ except AttributeError:
+ logger.info("Could not log the number of model parameters in SwanLab due to an AttributeError.")
+
+ # log the initial model architecture to an artifact
+ if self._log_model is not None:
+ logger.warning(
+ "SwanLab does not currently support the save mode functionality. "
+ "This feature will be available in a future release."
+ )
+ badge_markdown = (
+ f'[
]({self._swanlab.get_run().public.cloud.exp_url})'
+ )
+
+ modelcard.AUTOGENERATED_TRAINER_COMMENT += f"\n{badge_markdown}"
+
+ def on_train_begin(self, args, state, control, model=None, **kwargs):
+ if not self._initialized:
+ self.setup(args, state, model, **kwargs)
+
+ def on_train_end(self, args, state, control, model=None, processing_class=None, **kwargs):
+ if self._log_model is not None and self._initialized and state.is_world_process_zero:
+ logger.warning(
+ "SwanLab does not currently support the save mode functionality. "
+ "This feature will be available in a future release."
+ )
+
+ def on_log(self, args, state, control, model=None, logs=None, **kwargs):
+ single_value_scalars = [
+ "train_runtime",
+ "train_samples_per_second",
+ "train_steps_per_second",
+ "train_loss",
+ "total_flos",
+ ]
+
+ if not self._initialized:
+ self.setup(args, state, model)
+ if state.is_world_process_zero:
+ for k, v in logs.items():
+ if k in single_value_scalars:
+ self._swanlab.log({f"single_value/{k}": v}, step=state.global_step)
+ non_scalar_logs = {k: v for k, v in logs.items() if k not in single_value_scalars}
+ non_scalar_logs = rewrite_logs(non_scalar_logs)
+ self._swanlab.log({**non_scalar_logs, "train/global_step": state.global_step}, step=state.global_step)
+
+ def on_save(self, args, state, control, **kwargs):
+ if self._log_model is not None and self._initialized and state.is_world_process_zero:
+ logger.warning(
+ "SwanLab does not currently support the save mode functionality. "
+ "This feature will be available in a future release."
+ )
+
+ def on_predict(self, args, state, control, metrics, **kwargs):
+ if not self._initialized:
+ self.setup(args, state, **kwargs)
+ if state.is_world_process_zero:
+ metrics = rewrite_logs(metrics)
+ self._swanlab.log(metrics)
+
+
+INTEGRATION_TO_CALLBACK = {
+ "azure_ml": AzureMLCallback,
+ "comet_ml": CometCallback,
+ "mlflow": MLflowCallback,
+ "neptune": NeptuneCallback,
+ "tensorboard": TensorBoardCallback,
+ "wandb": WandbCallback,
+ "codecarbon": CodeCarbonCallback,
+ "clearml": ClearMLCallback,
+ "dagshub": DagsHubCallback,
+ "flyte": FlyteCallback,
+ "dvclive": DVCLiveCallback,
+ "swanlab": SwanLabCallback,
+}
+
+
+def get_reporting_integration_callbacks(report_to):
+ if report_to is None:
+ return []
+
+ if isinstance(report_to, str):
+ if "none" == report_to:
+ return []
+ elif "all" == report_to:
+ report_to = get_available_reporting_integrations()
+ else:
+ report_to = [report_to]
+
+ for integration in report_to:
+ if integration not in INTEGRATION_TO_CALLBACK:
+ raise ValueError(
+ f"{integration} is not supported, only {', '.join(INTEGRATION_TO_CALLBACK.keys())} are supported."
+ )
+
+ return [INTEGRATION_TO_CALLBACK[integration] for integration in report_to]
diff --git a/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/logging.py b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/logging.py
new file mode 100644
index 0000000000000000000000000000000000000000..a2915e167a0bcf9d61aeaca269608b98320c3c54
--- /dev/null
+++ b/qian_sg_backup/qwen3_hyp/qwen3_hyp_p047m_t1_8_suite-owt_scaling_v3_family-qwen3_geometry_id-hyp_variant-base_init_slope-std/attempt3_20260310_154910/logging.py
@@ -0,0 +1,410 @@
+# Copyright 2020 Optuna, Hugging Face
+#
+# 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.
+"""Logging utilities."""
+
+import functools
+import logging
+import os
+import sys
+import threading
+from logging import (
+ CRITICAL, # NOQA
+ DEBUG, # NOQA
+ ERROR, # NOQA
+ FATAL, # NOQA
+ INFO, # NOQA
+ NOTSET, # NOQA
+ WARN, # NOQA
+ WARNING, # NOQA
+)
+from logging import captureWarnings as _captureWarnings
+from typing import Optional
+
+import huggingface_hub.utils as hf_hub_utils
+from tqdm import auto as tqdm_lib
+
+
+_lock = threading.Lock()
+_default_handler: Optional[logging.Handler] = None
+
+log_levels = {
+ "detail": logging.DEBUG, # will also print filename and line number
+ "debug": logging.DEBUG,
+ "info": logging.INFO,
+ "warning": logging.WARNING,
+ "error": logging.ERROR,
+ "critical": logging.CRITICAL,
+}
+
+_default_log_level = logging.WARNING
+
+_tqdm_active = not hf_hub_utils.are_progress_bars_disabled()
+
+
+def _get_default_logging_level():
+ """
+ If TRANSFORMERS_VERBOSITY env var is set to one of the valid choices return that as the new default level. If it is
+ not - fall back to `_default_log_level`
+ """
+ env_level_str = os.getenv("TRANSFORMERS_VERBOSITY", None)
+ if env_level_str:
+ if env_level_str in log_levels:
+ return log_levels[env_level_str]
+ else:
+ logging.getLogger().warning(
+ f"Unknown option TRANSFORMERS_VERBOSITY={env_level_str}, "
+ f"has to be one of: {', '.join(log_levels.keys())}"
+ )
+ return _default_log_level
+
+
+def _get_library_name() -> str:
+ return __name__.split(".")[0]
+
+
+def _get_library_root_logger() -> logging.Logger:
+ return logging.getLogger(_get_library_name())
+
+
+def _configure_library_root_logger() -> None:
+ global _default_handler
+
+ with _lock:
+ if _default_handler:
+ # This library has already configured the library root logger.
+ return
+ _default_handler = logging.StreamHandler() # Set sys.stderr as stream.
+ # set defaults based on https://github.com/pyinstaller/pyinstaller/issues/7334#issuecomment-1357447176
+ if sys.stderr is None:
+ sys.stderr = open(os.devnull, "w")
+
+ _default_handler.flush = sys.stderr.flush
+
+ # Apply our default configuration to the library root logger.
+ library_root_logger = _get_library_root_logger()
+ library_root_logger.addHandler(_default_handler)
+ library_root_logger.setLevel(_get_default_logging_level())
+ # if logging level is debug, we add pathname and lineno to formatter for easy debugging
+ if os.getenv("TRANSFORMERS_VERBOSITY", None) == "detail":
+ formatter = logging.Formatter("[%(levelname)s|%(pathname)s:%(lineno)s] %(asctime)s >> %(message)s")
+ _default_handler.setFormatter(formatter)
+
+ is_ci = os.getenv("CI") is not None and os.getenv("CI").upper() in {"1", "ON", "YES", "TRUE"}
+ library_root_logger.propagate = True if is_ci else False
+
+
+def _reset_library_root_logger() -> None:
+ global _default_handler
+
+ with _lock:
+ if not _default_handler:
+ return
+
+ library_root_logger = _get_library_root_logger()
+ library_root_logger.removeHandler(_default_handler)
+ library_root_logger.setLevel(logging.NOTSET)
+ _default_handler = None
+
+
+def get_log_levels_dict():
+ return log_levels
+
+
+def captureWarnings(capture):
+ """
+ Calls the `captureWarnings` method from the logging library to enable management of the warnings emitted by the
+ `warnings` library.
+
+ Read more about this method here:
+ https://docs.python.org/3/library/logging.html#integration-with-the-warnings-module
+
+ All warnings will be logged through the `py.warnings` logger.
+
+ Careful: this method also adds a handler to this logger if it does not already have one, and updates the logging
+ level of that logger to the library's root logger.
+ """
+ logger = get_logger("py.warnings")
+
+ if not logger.handlers:
+ logger.addHandler(_default_handler)
+
+ logger.setLevel(_get_library_root_logger().level)
+
+ _captureWarnings(capture)
+
+
+def get_logger(name: Optional[str] = None) -> logging.Logger:
+ """
+ Return a logger with the specified name.
+
+ This function is not supposed to be directly accessed unless you are writing a custom transformers module.
+ """
+
+ if name is None:
+ name = _get_library_name()
+
+ _configure_library_root_logger()
+ return logging.getLogger(name)
+
+
+def get_verbosity() -> int:
+ """
+ Return the current level for the 🤗 Transformers's root logger as an int.
+
+ Returns:
+ `int`: The logging level.
+
+