Repository navigation
Conversation
Contributor
Author
This was referenced Sep 15, 2026
RyanPCo
force-pushed
the
ryanco/text-encoder-fixes
branch
from
September 16, 2026 17:30
5b608a7 to
35a858c
Compare
RyanPCo
force-pushed
the
ryanco/text-feature-cache
branch
from
September 16, 2026 17:30
3283194 to
09cfe20
Compare
This was referenced Sep 16, 2026
RyanPCo
force-pushed
the
ryanco/text-feature-cache
branch
from
September 17, 2026 14:52
09cfe20 to
a8b6490
Compare
This was referenced Sep 17, 2026
…split Three-operator twin of mecka_fold_flagship_opsplit_hpt_6d (and the wider sibling of the single-operator topop config): train on the 3 largest flagship fold-clothes operators (1,079 eps / 206,276 frames, 65.3 % of flagship), none of them held out. Seen-op val is the seeded 5 % episode split within those operators (53 eps / 9,826 frames); the unseen-operator head is inherited unchanged, so the runs stay comparable. Lane A keys measured from the real split on a compute node 2026-09-15: metric_frames_per_episode valid=floor(5120/53)=96, train_viz=floor(5120/1026)=4 (single-rank values; trainHydra scales them by the world size) (unseen 26 inherited); video_episodes pins one seen-val episode per train operator (191+223+226 = 640 f = exactly the 10-batch video budget) and inherits the train_viz and unseen pins. Co-Authored-By: Claude Opus 5 (1M context) <[email protected]> Claude-Session: https://claude.ai/code/session_01VBmhpVy7ZcMtsBJWh38cJp
…only grad-norm MAD Validating one checkpoint twice now gives identical prompts, actions and metrics. * HPT gets a per-pass RNG seed, base + rank * stride + pass counter, where the counter advances once per forward_eval call and is rewound by ModelWrapper.on_validation_start. forward_eval builds a torch.Generator from it and threads it down HPTModel.forward -> DenoisingPolicy.forward -> sample_action, and forks/seeds the global RNG around the BC val-loss call (DenoisingPolicy.predict draws its own noise and timestep, out of reach of any generator kwarg). Every generator kwarg defaults to None, so training and all other callers are untouched. The val-loss seed is offset by EVAL_LOSS_STRIDE rather than by 1: a torch.Generator seeded with S draws the same stream as a global manual_seed(S), so +1 would have handed batch i's val loss exactly the noise batch i+1 samples with. eval_base_seed is a real __init__ kwarg, so a second eval pass can be made to draw different noise on purpose. * _build_prompts picks annotations with that same seeded RNG whenever the algo is in eval mode, regardless of annotation_sampling_mode, so eval keeps prompt variety instead of collapsing to sample[0]. Train mode is unchanged. * The MAD grad-norm spike detector in on_after_backward is now log-only: it still logs the raw norm, the threshold, the flag and the [GRAD_NORM_SPIKE] line, but never rescales gradients, and every step (flagged or not) goes into the rolling median/MAD window. Fixed clipping moves to Lightning: trainer/default.yaml gains gradient_clip_val (null for now) and gradient_clip_algorithm: norm. The on_before_optimizer_step hook goes with it: Precision._after_closure calls that hook BEFORE _clip_gradients, so "policy_grad_norms_clipped" was the pre-clip norm under another name -- bit-identical to policy_grad_norms_raw, and no more informative once gradient_clip_val is set. Co-Authored-By: Claude Opus 5 (1M context) <[email protected]> Claude-Session: https://claude.ai/code/session_01VBmhpVy7ZcMtsBJWh38cJp
torchvision draws one random parameter set per call, so applying the configured train_image_augs to the stacked (B, 3, H, W) camera batch gave every image in the batch the same jitter. Wrap train_image_augs in a new PerSampleAugs module so each image draws its own parameters, and move the inline aug block of _robomimic_to_hpt_data into _apply_image_augs with the same gating. Deterministic eval augs stay batched. _apply_image_augs also checks train_image_augs before calling it: __init__ accepts None and train mode would have raised TypeError. The loop is not free: measured on an A40 at batch 64 with the shipped ColorJitter + Normalize, 13.8 -> 92.2 ms at 224x224 and 62.2 -> 105.5 ms at the camera's 360x640, per camera per training step. Correctness first here; the cost is removed further up the stack, where the jitter itself is vectorised over the batch. Co-Authored-By: Claude Fable 5.1 <[email protected]> Co-Authored-By: Claude Opus 5 (1M context) <[email protected]> Claude-Session: https://claude.ai/code/session_01VBmhpVy7ZcMtsBJWh38cJp
… LR group Two defects that would confound any backbone comparison: * a frozen `ResNet` backbone kept its BatchNorm layers in train mode, so it normalized with batch statistics and kept updating its running stats. `requires_grad = False` does not touch either. `train()` now puts a frozen backbone in eval mode. Latent today -- no shipped config freezes the ResNet -- so no run to date was affected. * the pretrained backbone and the randomly initialised trunk shared one LR. The new `model.backbone_lr_scale` key gives the encoders' backbones (`backbone_parameters()`) their own param group at `lr * scale`; at the default 1.0 the optimizer is built exactly as before, as a single group over `trainer.model.parameters()`. Nothing shipped sets it to anything else, so the two-group path is exercised only by its tests for now. Co-Authored-By: Claude Fable 5.1 <[email protected]> Co-Authored-By: Claude Opus 5 (1M context) <[email protected]> Claude-Session: https://claude.ai/code/session_01VBmhpVy7ZcMtsBJWh38cJp
ZarrDataset emits a real action_pad_mask (1.0 per real frame, 0.0 per repeat-last pad, honouring the annotation-cutoff EOS clamp) next to every horizoned action chunk. A new InterpolatePadMask transform carries it through the 30 -> 100 upsampling on the interpolators' own grid: an output step is real only if both source frames it interpolates between are real. HPT gains a use_pad_mask algo key (default false, declared in the two flagship model yamls; every other HPT config takes the __init__ default). With it on, process_batch_for_training builds pad_mask (B, S, 1) from the dataset mask instead of ones; DenoisingPolicy.loss_fn takes the mask and means the squared error over the valid steps only. With it off nothing changes: the mask is dropped from the batch and pad_mask stays all ones, which is exactly F.mse_loss. It ships off to match openpi, which asks LeRobotDataset for action_is_pad and then ignores it (Pi0.compute_loss means over the whole horizon). LeRobot's own policies do mask, but reduce with .mean() over all elements including the zeroed ones, so their loss shrinks as padding grows; this divides by the valid count instead, which keeps the gradient magnitude constant. Measured on the three pinned flagship episodes, 6.4-7.6 % of supervised action steps are repeat-last padding and 12.8-15.2 % of samples contain some, so turning it on changes the loss: a run before and a run after are not comparable, and the layout-aware eval metrics would have to be masked in the same change or train and val would be measuring different things. Padding is still repeat-last and no tail samples are dropped. Co-Authored-By: Claude Fable 5.1 <[email protected]> Co-Authored-By: Claude Opus 5 (1M context) <[email protected]> Claude-Session: https://claude.ai/code/session_01VBmhpVy7ZcMtsBJWh38cJp
…mbedding) The proprio keys the model consumes read K steps instead of one. Every change is gated on K > 1 (data) or history_len > 1 (model), and K = 1 / stride 1 is bit-for-bit today's pipeline. Data: Human.get_keymap gains proprio_history and history_stride; for K > 1 it marks the proprio entries with "history": K plus the stride, and ZarrDataset.__getitem__ reads the window idx - s*(K-1) ... idx in steps of s, clamped at the episode start and front-padded by repeating the first real frame, so the key is always (K, D) with the current frame last. A proprio_history_mask (1 = real, 0 = front padding) travels with it. The stride is why this ships on rather than at K = 1. Measured on a real flagship episode (30 fps, fold_clothes), the median wrist translation is 6.8 mm over 1 frame, 25.7 mm over 3, 42.9 mm over 5 and 80.8 mm over 10: an unstrided K = 3 window is 67 ms of past, close enough to a duplicate of the current pose to be mostly sensor noise, while the action chunk it conditions spans a full second. The flagship recipe therefore ships proprio_history: 3 / history_stride: 5 -- 0.33 s of past -- with the stem at history_len: 3. That changes the observation, so runs before and after this commit are not comparable, and history_dropout: 0.2 stops being inert: it is now a real hyperparameter, and the mitigation for the copycat/causal-confusion risk that proprio history classically brings. Transforms: the cartesian builders already treat a leading axis as batch; ActionChunkCoordinateFrameTransform / PoseCoordinateFrameTransform gain an opt-in target_history so a proprio target frame collapses to the CURRENT step (the action chunk must not move when history is added), and Reshape / QuaternionPoseToYPR tolerate the extra axis. observations.state.ee_pose is (K, 20). Eval: every human revert transform list starts with a SelectCurrentStep on its proprio key, so the val metrics, the cam-frame revert and the overlay video see the single current pose they always saw. No-op at K = 1 and idempotent. Norm stats are the baseline's: per-channel (D,) stats broadcast over (K, D), and both the training and the precompute stats pass force proprio_history=1 so K > 1 and K = 1 runs share one set. Model: MLPPolicyStem gains history_len / history_dropout. For history_len > 1 it owns a zero-init learned time embedding added after the projection, and HPT.stem_process skips the input-side sinusoid for such a stem (a constant [0, 1, 0, 1, ...] on normalized values is amplitude-1 noise on [-1, 1] data). history_dropout replaces the whole past by the current step per sample in training, exactly what an episode-start sample looks like after front padding. The K = 1 baseline stays one override away (proprio_history=1, history_stride=1, history_len=1) and the tests pin that it is bit-identical. Every other shipped config is untouched at K = 1, including the keypoint arm, whose wrist-frame proprio would need per-step framing before it can carry history. A compose test pins every train_*.yaml recipe's data K against its stem's history_len, because a recipe that inherits one model's data config and overrides the model is exactly how the two drift apart, and nothing else composes both halves. _stem_history_len takes the domain rather than matching a stem name by suffix: a cotrain run registers one state_ee_pose stem per domain, and a suffix match returns whichever comes first, so the human side was checked against eva's history_len. Latent while everything shipped K = 1; a blocker now. Co-Authored-By: Claude Fable 5.1 <[email protected]> Co-Authored-By: Claude Opus 5 (1M context) <[email protected]> Claude-Session: https://claude.ai/code/session_01VBmhpVy7ZcMtsBJWh38cJp
… padding `HPTModel.finalize_modules` ran `self.apply(self._init_weights)` over the whole model, which xavier-randomised the pretrained HF encoders inside the text stems (the frozen Qwen3-Embedding weights were overwritten in place and their biases zeroed). Modules that own pretrained weights now declare them through the `PretrainedWeights` mixin, and the init pass walks the tree without ever entering such a subtree, so the ResNet backbone, `T5Encoder` and any future pretrained stem are covered by the same mechanism. `HPT.__init__` also called `finalize_modules` twice; the first call is gone. Once the model is built, moved and upcast, `verify_pretrained_weights` hashes every declared pretrained submodule and compares it with the checkpoint read straight off the local HF snapshot, failing the run on a mismatch. It costs a second full copy of the encoder (~1.1 GB fp16 at 0.6B) and ~3 s of hashing per rank, so `verify_pretrained: false` turns it off; it is on by default, because a silently re-initialised encoder trains to a plausible wrong result. Only stems that own an HF snapshot are checked -- `ResNet` returns no reference, so it is protected from the init pass but not hash-verified, which the mixin docstring now says. Finally, `QwenPerTokenEncoder.compute_latent` now passes the tokenizer attention mask into its cross-attention. Padded positions are zeroed before the biased `proj`, so they were the learned constant `proj.bias`, not zero, and got real attention weight: text-conditioned results from before this commit are not comparable with results after it. Co-Authored-By: Claude Opus 5 (1M context) <[email protected]> Claude-Session: https://claude.ai/code/session_01VBmhpVy7ZcMtsBJWh38cJp
The annotation vocabulary is tiny next to the number of training steps, so a frozen text encoder re-runs the same 600M-parameter forward every step. `_Qwen3BaseEncoder` now keeps a lazy in-memory cache of the per-token hidden states keyed by the exact prompt string: misses go through the tokenizer + encoder in ONE sub-batch of unique strings, their un-padded (L_i, D) rows are stored detached, and the batch is re-assembled with the same LEFT padding the tokenizer produces, so `forward_with_mask` / the pooled path are unchanged. Only legal when the encoder is frozen - a trainable encoder must stay online. Knobs on the stem: `cache_text_features` (default true) and `text_cache_max_entries` (default 8192, LRU); both added to the model yamls with a Qwen stem. The cap counts entries, and an entry is the prompt's un-padded (L, D) states -- 4-25 KB for real annotations, 262 KB at max_length 128 -- so the default is ~100-200 MB of device memory in practice and 2 GB in the worst case; the docstring says so next to the knob. A batch carrying more unique prompts than the cap reads its rows from the miss pass rather than from the cache, because the LRU trim runs first and would otherwise evict entries the very next line looks up. `_apply` (to/cuda/float) clears the cache, and `clear_text_feature_cache()` / `text_feature_cache_stats()` are exposed. Padded positions come back as zeros instead of the encoder's pad-token output; every consumer masks them out. Cached-vs-online on 50 real prompts (fp16, each warmed alone: worst-case padding mismatch): max |diff| 3.6e-3 on the latents and 9.4e-2 on raw hidden states, exactly equal to the online path's own batch-padding noise (9.4e-2); in fp32 the same comparison is 1.1e-6 / 1.6e-4. Co-Authored-By: Claude Fable 5.1 <[email protected]> Co-Authored-By: Claude Opus 5 (1M context) <[email protected]> Claude-Session: https://claude.ai/code/session_01VBmhpVy7ZcMtsBJWh38cJp
aidang3019
force-pushed
the
ryanco/text-encoder-fixes
branch
from
September 30, 2026 01:07
f5d1a3a to
9d564a3
Compare
aidang3019
force-pushed
the
ryanco/text-feature-cache
branch
from
September 30, 2026 01:07
31b041b to
004383c
Compare
RyanPCo
force-pushed
the
ryanco/text-encoder-fixes
branch
from
October 12, 2026 20:26
9d564a3 to
7768bf6
Compare
Contributor
Author
|
Folded into #645 to keep the stack under Graphite's 50-branch cap; the changes are unchanged there. |
This was referenced Oct 12, 2026
Open
This was referenced Oct 22, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The annotation vocabulary is tiny next to the number of training steps, so a
frozen text encoder re-runs the same 600M-parameter forward every step.
_Qwen3BaseEncodernow keeps a lazy in-memory cache of the per-token hiddenstates keyed by the exact prompt string: misses go through the tokenizer +
encoder in ONE sub-batch of unique strings, their un-padded (L_i, D) rows are
stored detached, and the batch is re-assembled with the same LEFT padding the
tokenizer produces, so
forward_with_mask/ the pooled path are unchanged.Only legal when the encoder is frozen - a trainable encoder must stay online.
Knobs on the stem:
cache_text_features(default true) andtext_cache_max_entries(default 8192, LRU); both added to the four modelyamls with a Qwen stem.
_apply(to/cuda/float) clears the cache, andclear_text_feature_cache()/text_feature_cache_stats()are exposed.Padded positions come back as zeros instead of the encoder's pad-token output;
every consumer masks them out. Cached-vs-online on 50 real prompts (fp16, each
warmed alone: worst-case padding mismatch): max |diff| 3.6e-3 on the latents
and 9.4e-2 on raw hidden states, exactly equal to the online path's own
batch-padding noise (9.4e-2); in fp32 the same comparison is 1.1e-6 / 1.6e-4.
Co-Authored-By: Claude Fable 5.1 [email protected]