mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-29 20:49:42 +00:00
refactor(pi052): remove redundant policy code
This commit is contained in:
Pepijn
@@ -1,227 +0,0 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2026 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.
|
||||
|
||||
"""Audit or backfill checkpoint-local FAST artifacts for PI052 model repositories."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import io
|
||||
import json
|
||||
from pathlib import Path, PurePosixPath
|
||||
from typing import Any
|
||||
|
||||
from huggingface_hub import CommitOperationAdd, HfApi, hf_hub_download
|
||||
|
||||
DEFAULT_REPOSITORIES = (
|
||||
"pepijn223/pi052_atomic4_01_baseline",
|
||||
"pepijn223/pi052_atomic4_02_lr_1e5",
|
||||
"pepijn223/pi052_atomic4_03_recipe_50_50",
|
||||
"pepijn223/pi052_atomic4_04_flow_weight_10",
|
||||
"pepijn223/pi052_atomic4_05_flow_repeat_1",
|
||||
"pepijn223/pi052_atomic4_06_ki_off",
|
||||
)
|
||||
CHECKPOINT_DIRECTORIES = (
|
||||
"",
|
||||
"checkpoints/003000/pretrained_model",
|
||||
"checkpoints/006000/pretrained_model",
|
||||
"checkpoints/009000/pretrained_model",
|
||||
"checkpoints/012000/pretrained_model",
|
||||
)
|
||||
TOKENIZER_DIRECTORY = "action_tokenizer"
|
||||
|
||||
|
||||
def artifact_fingerprint(files: list[tuple[str, bytes]]) -> str:
|
||||
digest = hashlib.sha256()
|
||||
for relative_path, content in sorted(files):
|
||||
encoded_path = relative_path.encode()
|
||||
digest.update(len(encoded_path).to_bytes(8, "big"))
|
||||
digest.update(encoded_path)
|
||||
digest.update(len(content).to_bytes(8, "big"))
|
||||
digest.update(content)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def tokenizer_files(tokenizer_path: Path) -> list[tuple[str, Path]]:
|
||||
return [
|
||||
(path.relative_to(tokenizer_path).as_posix(), path)
|
||||
for path in sorted(tokenizer_path.rglob("*"))
|
||||
if path.is_file()
|
||||
]
|
||||
|
||||
|
||||
def _repo_path(directory: str, filename: str) -> str:
|
||||
return (PurePosixPath(directory) / filename).as_posix() if directory else filename
|
||||
|
||||
|
||||
def _download_json(repo_id: str, path_in_repo: str, revision: str | None = None) -> dict[str, Any]:
|
||||
path = hf_hub_download(repo_id, path_in_repo, repo_type="model", revision=revision)
|
||||
return json.loads(Path(path).read_text())
|
||||
|
||||
|
||||
def make_portable_preprocessor(config: dict[str, Any]) -> dict[str, Any]:
|
||||
config = json.loads(json.dumps(config))
|
||||
action_steps = [
|
||||
step for step in config["steps"] if step.get("registry_name") == "action_tokenizer_processor"
|
||||
]
|
||||
if len(action_steps) != 1:
|
||||
raise ValueError(f"Expected one action tokenizer step, found {len(action_steps)}")
|
||||
action_step = action_steps[0]
|
||||
action_step["config"]["action_tokenizer_name"] = TOKENIZER_DIRECTORY
|
||||
action_step["artifacts"] = {"action_tokenizer_name": TOKENIZER_DIRECTORY}
|
||||
|
||||
recipe_steps = [
|
||||
step for step in config["steps"] if step.get("registry_name") == "render_messages_processor"
|
||||
]
|
||||
if len(recipe_steps) != 1 or not recipe_steps[0].get("config", {}).get("recipe"):
|
||||
raise ValueError("PI052 preprocessor does not contain an embedded training recipe")
|
||||
return config
|
||||
|
||||
|
||||
def _json_operation(path_in_repo: str, content: dict[str, Any]) -> CommitOperationAdd:
|
||||
serialized = (json.dumps(content, indent=2) + "\n").encode()
|
||||
return CommitOperationAdd(path_in_repo=path_in_repo, path_or_fileobj=io.BytesIO(serialized))
|
||||
|
||||
|
||||
def prepare_operations(
|
||||
repo_id: str,
|
||||
tokenizer_path: Path,
|
||||
revision: str | None = None,
|
||||
) -> list[CommitOperationAdd]:
|
||||
operations: list[CommitOperationAdd] = []
|
||||
files = tokenizer_files(tokenizer_path)
|
||||
for directory in CHECKPOINT_DIRECTORIES:
|
||||
preprocessor_path = _repo_path(directory, "policy_preprocessor.json")
|
||||
operations.append(
|
||||
_json_operation(
|
||||
preprocessor_path,
|
||||
make_portable_preprocessor(_download_json(repo_id, preprocessor_path, revision)),
|
||||
)
|
||||
)
|
||||
for relative_path, local_path in files:
|
||||
operations.append(
|
||||
CommitOperationAdd(
|
||||
path_in_repo=_repo_path(
|
||||
directory,
|
||||
f"{TOKENIZER_DIRECTORY}/{relative_path}",
|
||||
),
|
||||
path_or_fileobj=str(local_path),
|
||||
)
|
||||
)
|
||||
return operations
|
||||
|
||||
|
||||
def audit_repository(
|
||||
api: HfApi,
|
||||
repo_id: str,
|
||||
expected_tokenizer_fingerprint: str,
|
||||
revision: str | None = None,
|
||||
) -> None:
|
||||
info = api.model_info(repo_id, revision=revision)
|
||||
repository_files = {sibling.rfilename for sibling in info.siblings or []}
|
||||
|
||||
for directory in CHECKPOINT_DIRECTORIES:
|
||||
preprocessor_path = _repo_path(directory, "policy_preprocessor.json")
|
||||
policy_config_path = _repo_path(directory, "config.json")
|
||||
postprocessor_path = _repo_path(directory, "policy_postprocessor.json")
|
||||
for required_path in (preprocessor_path, policy_config_path, postprocessor_path):
|
||||
if required_path not in repository_files:
|
||||
raise FileNotFoundError(f"{repo_id}@{revision or 'main'} is missing {required_path}")
|
||||
|
||||
preprocessor = _download_json(repo_id, preprocessor_path, revision)
|
||||
portable_preprocessor = make_portable_preprocessor(preprocessor)
|
||||
if preprocessor != portable_preprocessor:
|
||||
raise ValueError(f"{repo_id}:{preprocessor_path} is not portable")
|
||||
|
||||
normalizer_steps = [
|
||||
step for step in preprocessor["steps"] if step.get("registry_name") == "normalizer_processor"
|
||||
]
|
||||
if len(normalizer_steps) != 1 or "state_file" not in normalizer_steps[0]:
|
||||
raise ValueError(f"{repo_id}:{preprocessor_path} is missing normalizer state metadata")
|
||||
normalizer_path = _repo_path(directory, normalizer_steps[0]["state_file"])
|
||||
if normalizer_path not in repository_files:
|
||||
raise FileNotFoundError(f"{repo_id} is missing {normalizer_path}")
|
||||
|
||||
remote_tokenizer_files: list[tuple[str, bytes]] = []
|
||||
for relative_path in _tokenizer_relative_paths(repository_files, directory):
|
||||
path_in_repo = _repo_path(directory, f"{TOKENIZER_DIRECTORY}/{relative_path}")
|
||||
downloaded = hf_hub_download(repo_id, path_in_repo, repo_type="model", revision=revision)
|
||||
remote_tokenizer_files.append((relative_path, Path(downloaded).read_bytes()))
|
||||
fingerprint = artifact_fingerprint(remote_tokenizer_files)
|
||||
if fingerprint != expected_tokenizer_fingerprint:
|
||||
raise ValueError(
|
||||
f"{repo_id}:{_repo_path(directory, TOKENIZER_DIRECTORY)} fingerprint "
|
||||
f"{fingerprint} != {expected_tokenizer_fingerprint}"
|
||||
)
|
||||
|
||||
|
||||
def _tokenizer_relative_paths(repository_files: set[str], directory: str) -> list[str]:
|
||||
prefix = _repo_path(directory, TOKENIZER_DIRECTORY).rstrip("/") + "/"
|
||||
paths = sorted(path.removeprefix(prefix) for path in repository_files if path.startswith(prefix))
|
||||
if not paths:
|
||||
raise FileNotFoundError(f"Missing tokenizer artifact directory {prefix.rstrip('/')}")
|
||||
return paths
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--tokenizer-path", type=Path, required=True)
|
||||
parser.add_argument("--repo-id", action="append", dest="repo_ids")
|
||||
parser.add_argument("--revision")
|
||||
parser.add_argument("--dry-run", action="store_true")
|
||||
parser.add_argument("--audit-only", action="store_true")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
tokenizer_path = args.tokenizer_path.resolve()
|
||||
if not tokenizer_path.is_dir():
|
||||
raise FileNotFoundError(f"Tokenizer directory does not exist: {tokenizer_path}")
|
||||
|
||||
files = tokenizer_files(tokenizer_path)
|
||||
fingerprint = artifact_fingerprint([(relative_path, path.read_bytes()) for relative_path, path in files])
|
||||
api = HfApi()
|
||||
repositories = tuple(args.repo_ids or DEFAULT_REPOSITORIES)
|
||||
print(f"Tokenizer fingerprint: {fingerprint}")
|
||||
|
||||
for repo_id in repositories:
|
||||
if args.audit_only:
|
||||
audit_repository(api, repo_id, fingerprint, args.revision)
|
||||
print(f"AUDIT OK {repo_id}@{args.revision or 'main'}")
|
||||
continue
|
||||
|
||||
operations = prepare_operations(repo_id, tokenizer_path, args.revision)
|
||||
if args.dry_run:
|
||||
print(f"DRY RUN {repo_id}: {len(operations)} files")
|
||||
for operation in operations:
|
||||
print(f" {operation.path_in_repo}")
|
||||
continue
|
||||
|
||||
commit = api.create_commit(
|
||||
repo_id=repo_id,
|
||||
repo_type="model",
|
||||
operations=operations,
|
||||
commit_message="Embed fitted FAST tokenizer for portable PI052 checkpoints",
|
||||
revision=args.revision,
|
||||
)
|
||||
audit_repository(api, repo_id, fingerprint, commit.oid)
|
||||
print(f"BACKFILLED {repo_id}@{commit.oid}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -17,7 +17,6 @@
|
||||
import builtins
|
||||
import json
|
||||
import logging
|
||||
import math
|
||||
from collections import deque
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Literal, TypedDict, Unpack
|
||||
@@ -31,7 +30,6 @@ from lerobot.utils.import_utils import _transformers_available, require_package
|
||||
|
||||
# Conditional import for type checking and lazy loading
|
||||
if TYPE_CHECKING or _transformers_available:
|
||||
from transformers.cache_utils import DynamicCache
|
||||
from transformers.models.auto import CONFIG_MAPPING
|
||||
from transformers.models.gemma import modeling_gemma
|
||||
from transformers.utils import cached_file
|
||||
@@ -44,7 +42,6 @@ if TYPE_CHECKING or _transformers_available:
|
||||
)
|
||||
else:
|
||||
CONFIG_MAPPING = None
|
||||
DynamicCache = None
|
||||
modeling_gemma = None
|
||||
PiGemmaForCausalLM = None
|
||||
_gated_residual = None
|
||||
@@ -60,6 +57,13 @@ from lerobot.utils.constants import (
|
||||
)
|
||||
|
||||
from ..common.flow_matching import sample_noise, sample_time_beta
|
||||
from ..common.vla_utils import (
|
||||
clone_past_key_values,
|
||||
create_sinusoidal_pos_embedding,
|
||||
make_att_2d_masks,
|
||||
pad_vector,
|
||||
resize_with_pad_torch,
|
||||
)
|
||||
from ..pretrained import PreTrainedPolicy, T
|
||||
from ..rtc.modeling_rtc import RTCProcessor
|
||||
from .configuration_pi05 import DEFAULT_IMAGE_SIZE, PI05Config
|
||||
@@ -149,173 +153,6 @@ def _load_weight_files(files: list[Path]) -> dict[str, Tensor]:
|
||||
return state_dict
|
||||
|
||||
|
||||
def get_safe_dtype(target_dtype, device_type):
|
||||
"""Get a safe dtype for the given device type."""
|
||||
if device_type == "mps" and target_dtype == torch.float64:
|
||||
return torch.float32
|
||||
if device_type == "cpu":
|
||||
# CPU doesn't support bfloat16, use float32 instead
|
||||
if target_dtype == torch.bfloat16:
|
||||
return torch.float32
|
||||
if target_dtype == torch.float64:
|
||||
return torch.float64
|
||||
return target_dtype
|
||||
|
||||
|
||||
def create_sinusoidal_pos_embedding( # see openpi `create_sinusoidal_pos_embedding` (exact copy)
|
||||
time: torch.Tensor, dimension: int, min_period: float, max_period: float, device="cpu"
|
||||
) -> Tensor:
|
||||
"""Computes sine-cosine positional embedding vectors for scalar positions."""
|
||||
if dimension % 2 != 0:
|
||||
raise ValueError(f"dimension ({dimension}) must be divisible by 2")
|
||||
|
||||
if time.ndim != 1:
|
||||
raise ValueError("The time tensor is expected to be of shape `(batch_size, )`.")
|
||||
|
||||
dtype = get_safe_dtype(torch.float64, device.type)
|
||||
fraction = torch.linspace(0.0, 1.0, dimension // 2, dtype=dtype, device=device)
|
||||
period = min_period * (max_period / min_period) ** fraction
|
||||
|
||||
# Compute the outer product
|
||||
scaling_factor = 1.0 / period * 2 * math.pi
|
||||
sin_input = scaling_factor[None, :] * time[:, None]
|
||||
return torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1)
|
||||
|
||||
|
||||
def sample_beta(alpha, beta, bsize, device): # see openpi `sample_beta` (exact copy)
|
||||
# Beta sampling uses _sample_dirichlet which isn't implemented for MPS, so sample on CPU
|
||||
alpha_t = torch.tensor(alpha, dtype=torch.float32)
|
||||
beta_t = torch.tensor(beta, dtype=torch.float32)
|
||||
dist = torch.distributions.Beta(alpha_t, beta_t)
|
||||
return dist.sample((bsize,)).to(device)
|
||||
|
||||
|
||||
def make_att_2d_masks(pad_masks, att_masks): # see openpi `make_att_2d_masks` (exact copy)
|
||||
"""Copied from big_vision.
|
||||
|
||||
Tokens can attend to valid inputs tokens which have a cumulative mask_ar
|
||||
smaller or equal to theirs. This way `mask_ar` int[B, N] can be used to
|
||||
setup several types of attention, for example:
|
||||
|
||||
[[1 1 1 1 1 1]]: pure causal attention.
|
||||
|
||||
[[0 0 0 1 1 1]]: prefix-lm attention. The first 3 tokens can attend between
|
||||
themselves and the last 3 tokens have a causal attention. The first
|
||||
entry could also be a 1 without changing behaviour.
|
||||
|
||||
[[1 0 1 0 1 0 0 1 0 0]]: causal attention between 4 blocks. Tokens of a
|
||||
block can attend all previous blocks and all tokens on the same block.
|
||||
|
||||
Args:
|
||||
input_mask: bool[B, N] true if its part of the input, false if padding.
|
||||
mask_ar: int32[B, N] mask that's 1 where previous tokens cannot depend on
|
||||
it and 0 where it shares the same attention mask as the previous token.
|
||||
"""
|
||||
if att_masks.ndim != 2:
|
||||
raise ValueError(att_masks.ndim)
|
||||
if pad_masks.ndim != 2:
|
||||
raise ValueError(pad_masks.ndim)
|
||||
|
||||
cumsum = torch.cumsum(att_masks, dim=1)
|
||||
att_2d_masks = cumsum[:, None, :] <= cumsum[:, :, None]
|
||||
pad_2d_masks = pad_masks[:, None, :] * pad_masks[:, :, None]
|
||||
return att_2d_masks & pad_2d_masks
|
||||
|
||||
|
||||
def clone_past_key_values(past_key_values):
|
||||
"""Clone the DynamicCache returned by prefix prefill for compiled denoising."""
|
||||
return DynamicCache(
|
||||
tuple(
|
||||
(keys.clone(), values.clone(), sliding_window) for keys, values, sliding_window in past_key_values
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def pad_vector(vector, new_dim):
|
||||
"""Pad the last dimension of a vector to new_dim with zeros.
|
||||
|
||||
Can be (batch_size x sequence_length x features_dimension)
|
||||
or (batch_size x features_dimension)
|
||||
"""
|
||||
if vector.shape[-1] >= new_dim:
|
||||
return vector
|
||||
return F.pad(vector, (0, new_dim - vector.shape[-1]))
|
||||
|
||||
|
||||
def resize_with_pad_torch( # see openpi `resize_with_pad_torch` (exact copy)
|
||||
images: torch.Tensor,
|
||||
height: int,
|
||||
width: int,
|
||||
mode: str = "bilinear",
|
||||
) -> torch.Tensor:
|
||||
"""PyTorch version of resize_with_pad. Resizes an image to a target height and width without distortion
|
||||
by padding with black. If the image is float32, it must be in the range [-1, 1].
|
||||
|
||||
Args:
|
||||
images: Tensor of shape [*b, h, w, c] or [*b, c, h, w]
|
||||
height: Target height
|
||||
width: Target width
|
||||
mode: Interpolation mode ('bilinear', 'nearest', etc.)
|
||||
|
||||
Returns:
|
||||
Resized and padded tensor with same shape format as input
|
||||
"""
|
||||
# Check if input is in channels-last format [*b, h, w, c] or channels-first [*b, c, h, w]
|
||||
if images.shape[-1] <= 4: # Assume channels-last format
|
||||
channels_last = True
|
||||
if images.dim() == 3:
|
||||
images = images.unsqueeze(0) # Add batch dimension
|
||||
images = images.permute(0, 3, 1, 2) # [b, h, w, c] -> [b, c, h, w]
|
||||
else:
|
||||
channels_last = False
|
||||
if images.dim() == 3:
|
||||
images = images.unsqueeze(0) # Add batch dimension
|
||||
|
||||
batch_size, channels, cur_height, cur_width = images.shape
|
||||
|
||||
# Calculate resize ratio
|
||||
ratio = max(cur_width / width, cur_height / height)
|
||||
resized_height = int(cur_height / ratio)
|
||||
resized_width = int(cur_width / ratio)
|
||||
|
||||
# Resize
|
||||
resized_images = F.interpolate(
|
||||
images,
|
||||
size=(resized_height, resized_width),
|
||||
mode=mode,
|
||||
align_corners=False if mode == "bilinear" else None,
|
||||
)
|
||||
|
||||
# Handle dtype-specific clipping
|
||||
if images.dtype == torch.uint8:
|
||||
resized_images = torch.round(resized_images).clamp(0, 255).to(torch.uint8)
|
||||
elif images.dtype == torch.float32:
|
||||
resized_images = resized_images.clamp(0.0, 1.0)
|
||||
else:
|
||||
raise ValueError(f"Unsupported image dtype: {images.dtype}")
|
||||
|
||||
# Calculate padding
|
||||
pad_h0, remainder_h = divmod(height - resized_height, 2)
|
||||
pad_h1 = pad_h0 + remainder_h
|
||||
pad_w0, remainder_w = divmod(width - resized_width, 2)
|
||||
pad_w1 = pad_w0 + remainder_w
|
||||
|
||||
# Pad
|
||||
constant_value = 0 if images.dtype == torch.uint8 else 0.0
|
||||
padded_images = F.pad(
|
||||
resized_images,
|
||||
(pad_w0, pad_w1, pad_h0, pad_h1), # left, right, top, bottom
|
||||
mode="constant",
|
||||
value=constant_value,
|
||||
)
|
||||
|
||||
# Convert back to original format if needed
|
||||
if channels_last:
|
||||
padded_images = padded_images.permute(0, 2, 3, 1) # [b, c, h, w] -> [b, h, w, c]
|
||||
|
||||
return padded_images
|
||||
|
||||
|
||||
# Define the complete layer computation function for gradient checkpointing
|
||||
def compute_layer_complete(inputs_embeds, attention_mask, position_ids, adarms_cond, layers, rotary_emb):
|
||||
query_states = []
|
||||
|
||||
@@ -970,7 +970,7 @@ class PI052Policy(PI05Policy):
|
||||
for p in norm.parameters():
|
||||
p.requires_grad_(True)
|
||||
layers = getattr(text_model, "layers", None)
|
||||
if isinstance(layers, (list, torch.nn.ModuleList)) and len(layers) > 0:
|
||||
if isinstance(layers, list | torch.nn.ModuleList) and len(layers) > 0:
|
||||
for p in layers[-1].parameters():
|
||||
p.requires_grad_(True)
|
||||
|
||||
@@ -1084,6 +1084,53 @@ class PI052Policy(PI05Policy):
|
||||
loss_dict["loss"] = total.detach().mean()
|
||||
return total, loss_dict
|
||||
|
||||
def _embed_supervised_prefix(
|
||||
self,
|
||||
batch: dict[str, Tensor],
|
||||
text_labels: Tensor | None,
|
||||
action_tokens: Tensor | None,
|
||||
action_mask: Tensor | None,
|
||||
*,
|
||||
suppress_prefix_grads: bool = False,
|
||||
) -> tuple[Tensor, Tensor, Tensor, int, int]:
|
||||
"""Embed images, language, and optional FAST supervision once."""
|
||||
images, img_masks = self._preprocess_images(batch)
|
||||
with torch.no_grad() if suppress_prefix_grads else nullcontext():
|
||||
prefix_embs, prefix_pad, prefix_att = self.model.embed_prefix(
|
||||
images,
|
||||
img_masks,
|
||||
batch[OBS_LANGUAGE_TOKENS],
|
||||
batch[OBS_LANGUAGE_ATTENTION_MASK],
|
||||
)
|
||||
non_fast_prefix_len = prefix_embs.shape[1]
|
||||
|
||||
if text_labels is not None:
|
||||
lang_start = non_fast_prefix_len - text_labels.shape[1]
|
||||
if lang_start >= 0:
|
||||
prefix_att = _mark_target_span_causal(
|
||||
prefix_att, text_labels, lang_start, non_fast_prefix_len
|
||||
)
|
||||
|
||||
fast_len = 0
|
||||
if action_tokens is not None and action_mask is not None:
|
||||
fast_emb = self.model.paligemma_with_expert.embed_language_tokens(action_tokens)
|
||||
fast_len = action_tokens.shape[1]
|
||||
prefix_embs = torch.cat([prefix_embs, fast_emb], dim=1)
|
||||
prefix_pad = torch.cat([prefix_pad, action_mask.to(prefix_pad.dtype)], dim=1)
|
||||
prefix_att = torch.cat(
|
||||
[
|
||||
prefix_att,
|
||||
torch.ones(
|
||||
(action_tokens.shape[0], fast_len),
|
||||
dtype=torch.bool,
|
||||
device=prefix_embs.device,
|
||||
),
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
|
||||
return prefix_embs, prefix_pad, prefix_att, non_fast_prefix_len, fast_len
|
||||
|
||||
def _compute_all_losses_fused(
|
||||
self,
|
||||
batch: dict[str, Tensor],
|
||||
@@ -1105,37 +1152,13 @@ class PI052Policy(PI05Policy):
|
||||
and getattr(self.config, "knowledge_insulation", False)
|
||||
)
|
||||
|
||||
# ---- prefix: images + language + (optional FAST) -------------
|
||||
images, img_masks = self._preprocess_images(batch)
|
||||
lang_tokens = batch[OBS_LANGUAGE_TOKENS]
|
||||
lang_masks = batch[OBS_LANGUAGE_ATTENTION_MASK]
|
||||
with torch.no_grad() if suppress_prefix_grads else nullcontext():
|
||||
prefix_embs, prefix_pad, prefix_att = self.model.embed_prefix(
|
||||
images, img_masks, lang_tokens, lang_masks
|
||||
)
|
||||
non_fast_prefix_len = prefix_embs.shape[1] # images + language only
|
||||
|
||||
# Make supervised text causal rather than a bidirectional copy task.
|
||||
if text_labels is not None:
|
||||
lang_start = non_fast_prefix_len - text_labels.shape[1]
|
||||
if lang_start >= 0:
|
||||
prefix_att = _mark_target_span_causal(
|
||||
prefix_att, text_labels, lang_start, non_fast_prefix_len
|
||||
)
|
||||
|
||||
fast_len = 0
|
||||
if action_tokens is not None and action_mask is not None:
|
||||
# Gemma embedding already applies its hidden-size scale.
|
||||
fast_emb = self.model.paligemma_with_expert.embed_language_tokens(action_tokens)
|
||||
fast_len = action_tokens.shape[1]
|
||||
ones_att = torch.ones(
|
||||
(action_tokens.shape[0], fast_len),
|
||||
dtype=torch.bool,
|
||||
device=prefix_embs.device,
|
||||
)
|
||||
prefix_embs = torch.cat([prefix_embs, fast_emb], dim=1)
|
||||
prefix_pad = torch.cat([prefix_pad, action_mask.to(prefix_pad.dtype)], dim=1)
|
||||
prefix_att = torch.cat([prefix_att, ones_att], dim=1)
|
||||
prefix_embs, prefix_pad, prefix_att, non_fast_prefix_len, fast_len = self._embed_supervised_prefix(
|
||||
batch,
|
||||
text_labels,
|
||||
action_tokens,
|
||||
action_mask,
|
||||
suppress_prefix_grads=suppress_prefix_grads,
|
||||
)
|
||||
|
||||
# Amortized flow reuses one VLM prefix across fresh denoising targets.
|
||||
num_repeats = int(getattr(self.config, "flow_num_repeats", 1))
|
||||
@@ -1447,90 +1470,34 @@ class PI052Policy(PI05Policy):
|
||||
the caller doesn't want that head.
|
||||
"""
|
||||
|
||||
images, img_masks = self._preprocess_images(batch)
|
||||
lang_tokens = batch[OBS_LANGUAGE_TOKENS]
|
||||
lang_masks = batch[OBS_LANGUAGE_ATTENTION_MASK]
|
||||
|
||||
prefix_embs, prefix_pad, prefix_att = self.model.embed_prefix(
|
||||
images, img_masks, lang_tokens, lang_masks
|
||||
prefix_embs, prefix_pad, prefix_att, _, fast_len = self._embed_supervised_prefix(
|
||||
batch,
|
||||
text_labels,
|
||||
action_tokens,
|
||||
action_mask,
|
||||
)
|
||||
|
||||
# Make supervised text causal before appending FAST tokens.
|
||||
if text_labels is not None:
|
||||
lang_start = prefix_embs.shape[1] - text_labels.shape[1]
|
||||
if lang_start >= 0:
|
||||
prefix_att = _mark_target_span_causal(
|
||||
prefix_att, text_labels, lang_start, prefix_embs.shape[1]
|
||||
)
|
||||
|
||||
fast_len = 0
|
||||
if action_tokens is not None and action_mask is not None:
|
||||
# embed_language_tokens already applies the Gemma sqrt(hidden) scale (tf>=5.4.0);
|
||||
# do not scale FAST action tokens again (would double-scale).
|
||||
fast_emb = self.model.paligemma_with_expert.embed_language_tokens(action_tokens)
|
||||
|
||||
fast_len = action_tokens.shape[1]
|
||||
ones_att = torch.ones(
|
||||
(action_tokens.shape[0], fast_len),
|
||||
dtype=torch.bool,
|
||||
device=prefix_embs.device,
|
||||
)
|
||||
full_embs = torch.cat([prefix_embs, fast_emb], dim=1)
|
||||
full_pad = torch.cat([prefix_pad, action_mask.to(prefix_pad.dtype)], dim=1)
|
||||
full_att = torch.cat([prefix_att, ones_att], dim=1)
|
||||
else:
|
||||
full_embs = prefix_embs
|
||||
full_pad = prefix_pad
|
||||
full_att = prefix_att
|
||||
|
||||
att_2d = make_att_2d_masks(full_pad, full_att)
|
||||
position_ids = torch.cumsum(full_pad, dim=1) - 1
|
||||
att_2d_4d = self.model._prepare_attention_masks_4d(att_2d, dtype=full_embs.dtype)
|
||||
att_2d = make_att_2d_masks(prefix_pad, prefix_att)
|
||||
position_ids = torch.cumsum(prefix_pad, dim=1) - 1
|
||||
att_2d_4d = self.model._prepare_attention_masks_4d(att_2d, dtype=prefix_embs.dtype)
|
||||
|
||||
(vlm_out, _), _ = self.model.paligemma_with_expert.forward(
|
||||
attention_mask=att_2d_4d,
|
||||
position_ids=position_ids,
|
||||
past_key_values=None,
|
||||
inputs_embeds=[full_embs, None],
|
||||
inputs_embeds=[prefix_embs, None],
|
||||
use_cache=False,
|
||||
)
|
||||
if vlm_out is None:
|
||||
raise RuntimeError("PI052 text+fast loss: VLM forward returned no hidden states.")
|
||||
|
||||
lm_head = self.model.paligemma_with_expert.paligemma.lm_head
|
||||
|
||||
text_loss: Tensor | None = None
|
||||
if text_labels is not None:
|
||||
lang_len = text_labels.shape[1]
|
||||
# embed_prefix lays out as [images, language]; with FAST
|
||||
# appended the full sequence is [images, language, FAST].
|
||||
if fast_len > 0:
|
||||
text_hidden = vlm_out[:, -(fast_len + lang_len) : -fast_len, :]
|
||||
else:
|
||||
text_hidden = vlm_out[:, -lang_len:, :]
|
||||
text_loss = _shifted_lin_ce(
|
||||
text_hidden,
|
||||
lm_head.weight,
|
||||
text_labels,
|
||||
z_loss_weight=getattr(self.config, "text_ce_z_loss_weight", 0.0),
|
||||
compiled=self.config.use_compiled_text_ce,
|
||||
reduction=reduction,
|
||||
)
|
||||
|
||||
fast_loss: Tensor | None = None
|
||||
if action_tokens is not None and action_code_mask is not None and fast_len > 0:
|
||||
fast_hidden = vlm_out[:, -fast_len:, :]
|
||||
fast_loss = _fast_lin_ce(
|
||||
fast_hidden,
|
||||
lm_head.weight,
|
||||
action_tokens,
|
||||
action_code_mask,
|
||||
predict_actions_t,
|
||||
compiled=self.config.use_compiled_text_ce,
|
||||
reduction=reduction,
|
||||
)
|
||||
|
||||
return text_loss, fast_loss
|
||||
return self._prefix_ce_losses(
|
||||
vlm_out,
|
||||
text_labels,
|
||||
action_tokens,
|
||||
action_code_mask,
|
||||
fast_len,
|
||||
predict_actions_t,
|
||||
reduction,
|
||||
)
|
||||
|
||||
def select_message(
|
||||
self,
|
||||
|
||||
@@ -14,7 +14,7 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Literal
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
@@ -22,19 +22,11 @@ from torch.nn import functional as F # noqa: N812
|
||||
|
||||
from lerobot.utils.import_utils import _transformers_available
|
||||
|
||||
# Default PaliGemma SigLIP input resolution. Mirrors
|
||||
# ``pi05.configuration_pi05.DEFAULT_IMAGE_SIZE``; duplicated as a plain constant
|
||||
# to avoid importing the pi05 package here (which would create an import cycle:
|
||||
# pi_gemma -> pi05.__init__ -> modeling_pi05 -> pi_gemma).
|
||||
DEFAULT_IMAGE_SIZE = 224
|
||||
|
||||
if TYPE_CHECKING or _transformers_available:
|
||||
from transformers.cache_utils import DynamicCache
|
||||
from transformers.masking_utils import create_causal_mask
|
||||
from transformers.modeling_layers import GradientCheckpointingLayer
|
||||
from transformers.modeling_outputs import BaseModelOutputWithPast
|
||||
from transformers.models.auto import CONFIG_MAPPING
|
||||
from transformers.models.gemma import modeling_gemma
|
||||
from transformers.models.gemma.modeling_gemma import (
|
||||
GemmaAttention,
|
||||
GemmaConfig,
|
||||
@@ -58,8 +50,6 @@ else:
|
||||
GradientCheckpointingLayer = None
|
||||
BaseModelOutputWithPast = None
|
||||
create_causal_mask = None
|
||||
CONFIG_MAPPING = None
|
||||
modeling_gemma = None
|
||||
|
||||
|
||||
def _gated_residual(
|
||||
@@ -425,332 +415,3 @@ def sdpa_attention_forward(
|
||||
scale=scaling,
|
||||
)
|
||||
return attn_output.transpose(1, 2).contiguous(), None
|
||||
|
||||
|
||||
# Define the complete layer computation function for gradient checkpointing
|
||||
def compute_layer_complete(
|
||||
layer_idx, inputs_embeds, attention_mask, position_ids, adarms_cond, paligemma, gemma_expert
|
||||
):
|
||||
models = [paligemma.model.language_model, gemma_expert.model]
|
||||
query_states = []
|
||||
key_states = []
|
||||
value_states = []
|
||||
gates = []
|
||||
for i, hidden_states in enumerate(inputs_embeds):
|
||||
layer = models[i].layers[layer_idx]
|
||||
hidden_states, gate = layernorm_forward(layer.input_layernorm, hidden_states, adarms_cond[i])
|
||||
gates.append(gate)
|
||||
input_shape = hidden_states.shape[:-1]
|
||||
hidden_shape = (*input_shape, -1, layer.self_attn.head_dim)
|
||||
query_state = layer.self_attn.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
|
||||
key_state = layer.self_attn.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
|
||||
value_state = layer.self_attn.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
|
||||
query_states.append(query_state)
|
||||
key_states.append(key_state)
|
||||
value_states.append(value_state)
|
||||
# Concatenate and process attention
|
||||
query_states = torch.cat(query_states, dim=2)
|
||||
key_states = torch.cat(key_states, dim=2)
|
||||
value_states = torch.cat(value_states, dim=2)
|
||||
dummy_tensor = torch.zeros(
|
||||
query_states.shape[0],
|
||||
query_states.shape[2],
|
||||
query_states.shape[-1],
|
||||
device=query_states.device,
|
||||
dtype=query_states.dtype,
|
||||
)
|
||||
cos, sin = paligemma.model.language_model.rotary_emb(dummy_tensor, position_ids)
|
||||
query_states, key_states = modeling_gemma.apply_rotary_pos_emb(
|
||||
query_states, key_states, cos, sin, unsqueeze_dim=1
|
||||
)
|
||||
batch_size = query_states.shape[0]
|
||||
scaling = paligemma.model.language_model.layers[layer_idx].self_attn.scaling
|
||||
att_output, _ = sdpa_attention_forward(
|
||||
paligemma.model.language_model.layers[layer_idx].self_attn,
|
||||
query_states,
|
||||
key_states,
|
||||
value_states,
|
||||
attention_mask,
|
||||
scaling,
|
||||
)
|
||||
# Get head_dim from the current layer, not from the model
|
||||
head_dim = paligemma.model.language_model.layers[layer_idx].self_attn.head_dim
|
||||
att_output = att_output.reshape(batch_size, -1, 1 * 8 * head_dim)
|
||||
# Process layer outputs
|
||||
outputs_embeds = []
|
||||
start_pos = 0
|
||||
for i, hidden_states in enumerate(inputs_embeds):
|
||||
layer = models[i].layers[layer_idx]
|
||||
end_pos = start_pos + hidden_states.shape[1]
|
||||
if att_output.dtype != layer.self_attn.o_proj.weight.dtype:
|
||||
att_output = att_output.to(layer.self_attn.o_proj.weight.dtype)
|
||||
out_emb = layer.self_attn.o_proj(att_output[:, start_pos:end_pos])
|
||||
# first residual
|
||||
out_emb = _gated_residual(hidden_states, out_emb, gates[i])
|
||||
after_first_residual = out_emb.clone()
|
||||
out_emb, gate = layernorm_forward(layer.post_attention_layernorm, out_emb, adarms_cond[i])
|
||||
# Convert to bfloat16 if the next layer (mlp) uses bfloat16
|
||||
if layer.mlp.up_proj.weight.dtype == torch.bfloat16:
|
||||
out_emb = out_emb.to(dtype=torch.bfloat16)
|
||||
out_emb = layer.mlp(out_emb)
|
||||
# second residual
|
||||
out_emb = _gated_residual(after_first_residual, out_emb, gate)
|
||||
outputs_embeds.append(out_emb)
|
||||
start_pos = end_pos
|
||||
return outputs_embeds
|
||||
|
||||
|
||||
class GemmaVariantConfig: # see openpi `gemma.py: Config`
|
||||
"""Configuration for Gemma model variants."""
|
||||
|
||||
def __init__(self, width, depth, mlp_dim, num_heads, num_kv_heads, head_dim):
|
||||
self.width = width
|
||||
self.depth = depth
|
||||
self.mlp_dim = mlp_dim
|
||||
self.num_heads = num_heads
|
||||
self.num_kv_heads = num_kv_heads
|
||||
self.head_dim = head_dim
|
||||
|
||||
|
||||
def get_gemma_config(variant: str) -> GemmaVariantConfig: # see openpi `gemma.py: get_config`
|
||||
"""Returns config for specified gemma variant."""
|
||||
if variant == "gemma_300m":
|
||||
return GemmaVariantConfig(
|
||||
width=1024,
|
||||
depth=18,
|
||||
mlp_dim=4096,
|
||||
num_heads=8,
|
||||
num_kv_heads=1,
|
||||
head_dim=256,
|
||||
)
|
||||
elif variant == "gemma_2b":
|
||||
return GemmaVariantConfig(
|
||||
width=2048,
|
||||
depth=18,
|
||||
mlp_dim=16_384,
|
||||
num_heads=8,
|
||||
num_kv_heads=1,
|
||||
head_dim=256,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unknown variant: {variant}")
|
||||
|
||||
|
||||
class PaliGemmaWithExpertModel(
|
||||
nn.Module
|
||||
): # see openpi `gemma_pytorch.py: PaliGemmaWithExpertModel` this class is almost a exact copy of PaliGemmaWithExpertModel in openpi
|
||||
"""PaliGemma model with action expert for PI05."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vlm_config,
|
||||
action_expert_config,
|
||||
use_adarms=None,
|
||||
precision: Literal["bfloat16", "float32"] = "bfloat16",
|
||||
image_size: int = DEFAULT_IMAGE_SIZE,
|
||||
freeze_vision_encoder: bool = False,
|
||||
train_expert_only: bool = False,
|
||||
):
|
||||
if use_adarms is None:
|
||||
use_adarms = [False, False]
|
||||
super().__init__()
|
||||
self.freeze_vision_encoder = freeze_vision_encoder
|
||||
self.train_expert_only = train_expert_only
|
||||
|
||||
vlm_config_hf = CONFIG_MAPPING["paligemma"]()
|
||||
vlm_config_hf._vocab_size = 257152 # noqa: SLF001
|
||||
vlm_config_hf.image_token_index = 257152
|
||||
vlm_config_hf.text_config.hidden_size = vlm_config.width
|
||||
vlm_config_hf.text_config.intermediate_size = vlm_config.mlp_dim
|
||||
vlm_config_hf.text_config.num_attention_heads = vlm_config.num_heads
|
||||
vlm_config_hf.text_config.head_dim = vlm_config.head_dim
|
||||
vlm_config_hf.text_config.num_hidden_layers = vlm_config.depth
|
||||
vlm_config_hf.text_config.num_key_value_heads = vlm_config.num_kv_heads
|
||||
vlm_config_hf.text_config.hidden_activation = "gelu_pytorch_tanh"
|
||||
vlm_config_hf.text_config.dtype = "float32"
|
||||
vlm_config_hf.text_config.vocab_size = 257152
|
||||
vlm_config_hf.text_config.use_adarms = use_adarms[0]
|
||||
vlm_config_hf.text_config.adarms_cond_dim = vlm_config.width if use_adarms[0] else None
|
||||
vlm_config_hf.vision_config.image_size = image_size
|
||||
vlm_config_hf.vision_config.intermediate_size = 4304
|
||||
vlm_config_hf.vision_config.projection_dim = 2048
|
||||
vlm_config_hf.vision_config.projector_hidden_act = "gelu_fast"
|
||||
vlm_config_hf.vision_config.dtype = "float32"
|
||||
|
||||
action_expert_config_hf = CONFIG_MAPPING["gemma"](
|
||||
head_dim=action_expert_config.head_dim,
|
||||
hidden_size=action_expert_config.width,
|
||||
intermediate_size=action_expert_config.mlp_dim,
|
||||
num_attention_heads=action_expert_config.num_heads,
|
||||
num_hidden_layers=action_expert_config.depth,
|
||||
num_key_value_heads=action_expert_config.num_kv_heads,
|
||||
vocab_size=257152,
|
||||
hidden_activation="gelu_pytorch_tanh",
|
||||
dtype="float32",
|
||||
use_adarms=use_adarms[1],
|
||||
adarms_cond_dim=action_expert_config.width if use_adarms[1] else None,
|
||||
)
|
||||
|
||||
self.paligemma = PaliGemmaForConditionalGenerationWithPiGemma(config=vlm_config_hf)
|
||||
self.gemma_expert = PiGemmaForCausalLM(config=action_expert_config_hf)
|
||||
self.gemma_expert.model.embed_tokens = None
|
||||
|
||||
self.to_bfloat16_for_selected_params(precision)
|
||||
self._set_requires_grad()
|
||||
|
||||
def to_bfloat16_for_selected_params(self, precision: Literal["bfloat16", "float32"] = "bfloat16"):
|
||||
if precision == "bfloat16":
|
||||
self.to(dtype=torch.bfloat16)
|
||||
elif precision == "float32":
|
||||
self.to(dtype=torch.float32)
|
||||
return
|
||||
else:
|
||||
raise ValueError(f"Invalid precision: {precision}")
|
||||
|
||||
# Keep full vision path in float32 so we never toggle (toggle causes optimizer
|
||||
# "same dtype" error). Saves memory vs full float32; more memory than only 3 params.
|
||||
params_to_keep_float32 = [
|
||||
"vision_tower",
|
||||
"multi_modal_projector",
|
||||
"lm_head",
|
||||
"input_layernorm",
|
||||
"post_attention_layernorm",
|
||||
"model.norm",
|
||||
]
|
||||
|
||||
for name, param in self.named_parameters():
|
||||
if any(selector in name for selector in params_to_keep_float32):
|
||||
param.data = param.data.to(dtype=torch.float32)
|
||||
|
||||
def _set_requires_grad(self):
|
||||
if self.freeze_vision_encoder:
|
||||
self.paligemma.model.vision_tower.eval()
|
||||
for param in self.paligemma.model.vision_tower.parameters():
|
||||
param.requires_grad = False
|
||||
if self.train_expert_only:
|
||||
self.paligemma.eval()
|
||||
for param in self.paligemma.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
def train(self, mode: bool = True):
|
||||
super().train(mode)
|
||||
if self.freeze_vision_encoder:
|
||||
self.paligemma.model.vision_tower.eval()
|
||||
if self.train_expert_only:
|
||||
self.paligemma.eval()
|
||||
|
||||
def embed_image(self, image: torch.Tensor):
|
||||
# Vision tower and multi_modal_projector are kept in float32 (params_to_keep_float32).
|
||||
out_dtype = image.dtype
|
||||
if image.dtype != torch.float32:
|
||||
image = image.to(torch.float32)
|
||||
image_outputs = self.paligemma.model.get_image_features(image)
|
||||
# OpenPI / big_vision convention: image (soft) tokens are NOT scaled by the
|
||||
# Gemma embedder normalizer (sqrt(hidden_size)) — only text tokens are. lerobot/pi05_base
|
||||
# was trained in this regime, so scaling image features here over-scales them ~45x and
|
||||
# breaks the pretrained vision-language alignment. Keep image features un-normalized.
|
||||
features = image_outputs.pooler_output
|
||||
if features.dtype != out_dtype:
|
||||
features = features.to(out_dtype)
|
||||
return features
|
||||
|
||||
def embed_language_tokens(self, tokens: torch.Tensor):
|
||||
return self.paligemma.model.language_model.embed_tokens(tokens)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
attention_mask: torch.Tensor | None = None,
|
||||
position_ids: torch.LongTensor | None = None,
|
||||
past_key_values: list[torch.FloatTensor] | None = None,
|
||||
inputs_embeds: list[torch.FloatTensor] | None = None,
|
||||
use_cache: bool | None = None,
|
||||
adarms_cond: list[torch.Tensor] | None = None,
|
||||
):
|
||||
if adarms_cond is None:
|
||||
adarms_cond = [None, None]
|
||||
if inputs_embeds[1] is None:
|
||||
prefix_output = self.paligemma.model.language_model.forward(
|
||||
inputs_embeds=inputs_embeds[0],
|
||||
attention_mask=attention_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_values=past_key_values,
|
||||
use_cache=use_cache,
|
||||
adarms_cond=adarms_cond[0] if adarms_cond is not None else None,
|
||||
)
|
||||
prefix_past_key_values = prefix_output.past_key_values
|
||||
prefix_output = prefix_output.last_hidden_state
|
||||
suffix_output = None
|
||||
elif inputs_embeds[0] is None:
|
||||
suffix_output = self.gemma_expert.model.forward(
|
||||
inputs_embeds=inputs_embeds[1],
|
||||
attention_mask=attention_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_values=past_key_values,
|
||||
use_cache=use_cache,
|
||||
adarms_cond=adarms_cond[1] if adarms_cond is not None else None,
|
||||
)
|
||||
suffix_output = suffix_output.last_hidden_state
|
||||
prefix_output = None
|
||||
prefix_past_key_values = None
|
||||
else:
|
||||
models = [self.paligemma.model.language_model, self.gemma_expert.model]
|
||||
num_layers = self.paligemma.config.text_config.num_hidden_layers
|
||||
|
||||
# Check if gradient checkpointing is enabled for any of the models
|
||||
use_gradient_checkpointing = (
|
||||
hasattr(self.gemma_expert.model, "gradient_checkpointing")
|
||||
and self.gemma_expert.model.gradient_checkpointing
|
||||
and self.training
|
||||
) or (hasattr(self, "gradient_checkpointing") and self.gradient_checkpointing and self.training)
|
||||
|
||||
# Process all layers with gradient checkpointing if enabled
|
||||
for layer_idx in range(num_layers):
|
||||
if use_gradient_checkpointing:
|
||||
inputs_embeds = torch.utils.checkpoint.checkpoint(
|
||||
compute_layer_complete,
|
||||
layer_idx,
|
||||
inputs_embeds,
|
||||
attention_mask,
|
||||
position_ids,
|
||||
adarms_cond,
|
||||
use_reentrant=False,
|
||||
preserve_rng_state=False,
|
||||
paligemma=self.paligemma,
|
||||
gemma_expert=self.gemma_expert,
|
||||
)
|
||||
else:
|
||||
inputs_embeds = compute_layer_complete(
|
||||
layer_idx,
|
||||
inputs_embeds,
|
||||
attention_mask,
|
||||
position_ids,
|
||||
adarms_cond,
|
||||
paligemma=self.paligemma,
|
||||
gemma_expert=self.gemma_expert,
|
||||
)
|
||||
|
||||
# final norm
|
||||
def compute_final_norms(inputs_embeds, adarms_cond):
|
||||
outputs_embeds = []
|
||||
for i, hidden_states in enumerate(inputs_embeds):
|
||||
out_emb, _ = layernorm_forward(models[i].norm, hidden_states, adarms_cond[i])
|
||||
outputs_embeds.append(out_emb)
|
||||
return outputs_embeds
|
||||
|
||||
# Apply gradient checkpointing to final norm if enabled
|
||||
if use_gradient_checkpointing:
|
||||
outputs_embeds = torch.utils.checkpoint.checkpoint(
|
||||
compute_final_norms,
|
||||
inputs_embeds,
|
||||
adarms_cond,
|
||||
use_reentrant=False,
|
||||
preserve_rng_state=False,
|
||||
)
|
||||
else:
|
||||
outputs_embeds = compute_final_norms(inputs_embeds, adarms_cond)
|
||||
|
||||
prefix_output = outputs_embeds[0]
|
||||
suffix_output = outputs_embeds[1]
|
||||
prefix_past_key_values = None
|
||||
|
||||
return [prefix_output, suffix_output], prefix_past_key_values
|
||||
|
||||
@@ -1,67 +0,0 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2026 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 scripts.backfill_pi052_action_tokenizer import (
|
||||
CHECKPOINT_DIRECTORIES,
|
||||
DEFAULT_REPOSITORIES,
|
||||
artifact_fingerprint,
|
||||
make_portable_preprocessor,
|
||||
)
|
||||
|
||||
|
||||
def test_atomic4_backfill_covers_every_repository_and_checkpoint():
|
||||
assert len(DEFAULT_REPOSITORIES) == 6
|
||||
assert CHECKPOINT_DIRECTORIES == (
|
||||
"",
|
||||
"checkpoints/003000/pretrained_model",
|
||||
"checkpoints/006000/pretrained_model",
|
||||
"checkpoints/009000/pretrained_model",
|
||||
"checkpoints/012000/pretrained_model",
|
||||
)
|
||||
|
||||
|
||||
def test_backfill_embeds_recipe_and_declares_relative_tokenizer():
|
||||
recipe = {"messages": [{"role": "user", "content": "${task}", "stream": "low_level"}]}
|
||||
preprocessor = {
|
||||
"name": "policy_preprocessor",
|
||||
"steps": [
|
||||
{
|
||||
"registry_name": "normalizer_processor",
|
||||
"config": {},
|
||||
"state_file": "normalizer.safetensors",
|
||||
},
|
||||
{"registry_name": "render_messages_processor", "config": {"recipe": recipe}},
|
||||
{
|
||||
"registry_name": "action_tokenizer_processor",
|
||||
"config": {"action_tokenizer_name": "/fsx/original/tokenizer"},
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
portable = make_portable_preprocessor(preprocessor)
|
||||
|
||||
assert portable["steps"][1]["config"]["recipe"] == recipe
|
||||
assert portable["steps"][2]["config"]["action_tokenizer_name"] == "action_tokenizer"
|
||||
assert portable["steps"][2]["artifacts"] == {"action_tokenizer_name": "action_tokenizer"}
|
||||
assert preprocessor["steps"][2]["config"]["action_tokenizer_name"].startswith("/fsx/")
|
||||
|
||||
|
||||
def test_artifact_fingerprint_includes_paths_and_contents():
|
||||
first = artifact_fingerprint([("a/file", b"same"), ("b/file", b"content")])
|
||||
|
||||
assert first == artifact_fingerprint([("b/file", b"content"), ("a/file", b"same")])
|
||||
assert first != artifact_fingerprint([("a/renamed", b"same"), ("b/file", b"content")])
|
||||
assert first != artifact_fingerprint([("a/file", b"changed"), ("b/file", b"content")])
|
||||
Reference in New Issue
Block a user