diff --git a/src/lerobot/policies/fastwam/configuration_fastwam.py b/src/lerobot/policies/fastwam/configuration_fastwam.py index 5907769fa..9e76a2eb1 100644 --- a/src/lerobot/policies/fastwam/configuration_fastwam.py +++ b/src/lerobot/policies/fastwam/configuration_fastwam.py @@ -202,6 +202,11 @@ class FastWAMConfig(PreTrainedConfig): base_model_id: str | None = FASTWAM_BASE_MODEL_ID tokenizer_max_len: int = 128 load_text_encoder: bool = True + # Device for the frozen ~11GB UMT5-XXL text encoder. `None` keeps it on the main + # policy `device` (default). Set to e.g. "cpu" to keep it off the GPU and save VRAM; + # prompts are then encoded on that device and the resulting embeddings moved to the + # policy device. Trades GPU memory for slower (CPU) text encoding. + text_encoder_device: str | None = None mot_checkpoint_mixed_attn: bool = False torch_dtype: str = "bfloat16" prompt_template: str = ( diff --git a/src/lerobot/policies/fastwam/modeling_fastwam.py b/src/lerobot/policies/fastwam/modeling_fastwam.py index 0563242b5..dbe39e243 100644 --- a/src/lerobot/policies/fastwam/modeling_fastwam.py +++ b/src/lerobot/policies/fastwam/modeling_fastwam.py @@ -150,6 +150,24 @@ class FastWAMPolicy(PreTrainedPolicy): def reset(self) -> None: self._action_queue: deque[Tensor] = deque([], maxlen=self.config.n_action_steps) + # Per-episode text-embedding cache (mirrors LingBot-VA's `_prompt_embeds`). The task + # is fixed for an episode, so the ~11GB UMT5 encoder runs once on the first chunk and + # the resulting context is reused for every subsequent chunk. Cleared here on reset so + # a new episode's (possibly different) task is re-encoded. Proprio is still appended + # fresh each chunk downstream, so only the text-only context is cached. + self._cached_prompt: Any = None + self._cached_context: Tensor | None = None + self._cached_context_mask: Tensor | None = None + + def _encode_prompt_cached(self, prompt: Any) -> tuple[Tensor, Tensor]: + """Encode `prompt` to `(context, context_mask)`, reusing the cache when the prompt is + unchanged so UMT5 runs at most once per episode (per distinct task).""" + if self._cached_context is None or self._cached_prompt != prompt: + context, context_mask = self.model.encode_prompt(prompt) + self._cached_prompt = prompt + self._cached_context = context + self._cached_context_mask = context_mask + return self._cached_context, self._cached_context_mask def _batch_to_training_sample(self, batch: dict[str, Tensor]) -> dict[str, Tensor]: """Adapt a standard LeRobot batch to the FastWAM-native sample that @@ -222,6 +240,15 @@ class FastWAMPolicy(PreTrainedPolicy): self.eval() infer_kwargs = _batch_to_infer_kwargs(batch=batch, config=self.config) + # Encode the task once per episode and reuse it (LingBot-VA parity): swap the raw + # `prompt` for the cached `context`/`context_mask` so `infer_action` skips `encode_prompt` + # and the text encoder isn't re-run every chunk. Skipped when the caller supplies its own + # precomputed `context` (the two are mutually exclusive downstream). + if infer_kwargs.get("context") is None and infer_kwargs.get("prompt") is not None: + context, context_mask = self._encode_prompt_cached(infer_kwargs["prompt"]) + infer_kwargs["prompt"] = None + infer_kwargs["context"] = context + infer_kwargs["context_mask"] = context_mask batch_size = _infer_kwargs_batch_size(infer_kwargs) if batch_size == 1: action = _action_from_model_output(self.model.infer_action(**infer_kwargs)) @@ -266,9 +293,10 @@ class FastWAMPolicy(PreTrainedPolicy): mixtures={"video": video_expert, "action": action_expert}, mot_checkpoint_mixed_attn=config.mot_checkpoint_mixed_attn, ) + text_encoder_device = config.text_encoder_device or device text_encoder = ( load_pretrained_wan_text_encoder( - model_id=config.text_encoder_model_id, torch_dtype=dtype, device=device + model_id=config.text_encoder_model_id, torch_dtype=dtype, device=text_encoder_device ) if config.load_text_encoder else None @@ -279,6 +307,7 @@ class FastWAMPolicy(PreTrainedPolicy): mot=mot, vae=load_pretrained_wan_vae(torch_dtype=dtype, device=device), text_encoder=text_encoder, + text_encoder_device=config.text_encoder_device, tokenizer=build_wan_tokenizer( model_id=config.tokenizer_model_id, tokenizer_max_len=config.tokenizer_max_len ), diff --git a/src/lerobot/policies/fastwam/wan/modular.py b/src/lerobot/policies/fastwam/wan/modular.py index add9c82ef..80cac4be8 100644 --- a/src/lerobot/policies/fastwam/wan/modular.py +++ b/src/lerobot/policies/fastwam/wan/modular.py @@ -839,6 +839,7 @@ class FastWAM(torch.nn.Module): text_dim: int | None = None, proprio_dim: int | None = None, device: str = "cpu", + text_encoder_device: str | torch.device | None = None, torch_dtype: torch.dtype = torch.float32, video_train_shift: float = 5.0, video_infer_shift: float = 5.0, @@ -908,11 +909,21 @@ class FastWAM(torch.nn.Module): self.infer_scheduler = self.infer_video_scheduler self.device = torch.device(device) + # When pinned (e.g. "cpu"), the frozen text encoder stays on this device instead + # of following the model onto the GPU — `_apply` skips it and `encode_prompt` runs + # it here, moving embeddings back to `self.device`. `None` = follow `self.device`. + self._text_encoder_device = ( + torch.device(text_encoder_device) if text_encoder_device is not None else None + ) self.torch_dtype = torch_dtype self.loss_lambda_video = float(loss_lambda_video) self.loss_lambda_action = float(loss_lambda_action) self.to(self.device) + # `self.to` above (via `_apply`) skips a pinned text encoder; make sure it actually + # sits on the pinned device (it was loaded there, but this is a cheap safety net). + if self.text_encoder is not None and self._text_encoder_device is not None: + self.text_encoder._apply(lambda t: t.to(self._text_encoder_device)) @classmethod def from_wan22_pretrained( @@ -1003,7 +1014,8 @@ class FastWAM(torch.nn.Module): # while staying out of `state_dict()` / `parameters()`. super()._apply(fn, *args, **kwargs) self.vae._apply(fn) - if self.text_encoder is not None: + # A pinned text encoder (e.g. on CPU) must NOT follow device moves — leave it put. + if self.text_encoder is not None and self._text_encoder_device is None: self.text_encoder._apply(fn) return self @@ -1024,9 +1036,12 @@ class FastWAM(torch.nn.Module): "Prompt encoding requires loaded text encoder/tokenizer. " "Set `load_text_encoder=true` or provide precomputed `context/context_mask`." ) + # Run the encoder on its own device (may be pinned to CPU to save VRAM), then + # move the resulting embeddings/mask to the model device for the DiT. + te_device = self._text_encoder_device or self.device ids, mask = self.tokenizer(prompt, return_mask=True, add_special_tokens=True) - ids = ids.to(self.device) - mask = mask.to(self.device, dtype=torch.bool) + ids = ids.to(te_device) + mask = mask.to(te_device, dtype=torch.bool) prompt_emb = self.text_encoder(ids, mask) seq_lens = mask.gt(0).sum(dim=1).long() for i, v in enumerate(seq_lens): @@ -1034,7 +1049,7 @@ class FastWAM(torch.nn.Module): # Match FastWAM/Wan2.2 context semantics: padding embeddings are zeroed, # while cross-attention still sees a fixed-length context. mask = torch.ones_like(mask) - return prompt_emb.to(device=self.device), mask + return prompt_emb.to(device=self.device), mask.to(device=self.device) def _append_proprio_to_context( self, diff --git a/src/lerobot/policies/lingbot_va/modeling_lingbot_va.py b/src/lerobot/policies/lingbot_va/modeling_lingbot_va.py index b1b82f026..72563f991 100644 --- a/src/lerobot/policies/lingbot_va/modeling_lingbot_va.py +++ b/src/lerobot/policies/lingbot_va/modeling_lingbot_va.py @@ -573,7 +573,6 @@ class LingBotVAPolicy(PreTrainedPolicy): if size is None: size = (self.config.height, self.config.width) img = F.interpolate(img, size=size, mode="bilinear", align_corners=False) - print(f"[lingbot_va] {key} pre-scale range: [{img.min().item():.4f}, {img.max().item():.4f}] (expect ~[0,1])") img = img * 2.0 - 1.0 return img.unsqueeze(2).to(self.dtype) # [1, C, F=1, H, W]