Add Wall-X serving and Turtle2 TCP WebSocket bridge
Pre-commit / pre-commit (push) Canceled after 0s
Pre-commit / pre-commit (push) Canceled after 0s
This commit is contained in:
@@ -321,6 +321,20 @@ class ActionModelMixMin:
|
||||
noisy_action_emb = noisy_action_emb.to(
|
||||
inputs_embeds.device, inputs_embeds.dtype
|
||||
)
|
||||
action_token_count = int(mask.sum().item())
|
||||
if noisy_action_emb.ndim != 3 or (
|
||||
action_token_count * inputs_embeds.shape[-1]
|
||||
!= noisy_action_emb.numel()
|
||||
):
|
||||
raise ValueError(
|
||||
"Flow action placeholder/embedding mismatch: "
|
||||
f"input_ids={tuple(input_ids.shape)}, "
|
||||
f"action_tokens={action_token_count}, "
|
||||
f"action_chunk={tuple(action_chunk.shape)}, "
|
||||
f"inputs_embeds={tuple(inputs_embeds.shape)}, "
|
||||
f"noisy_action_emb={tuple(noisy_action_emb.shape)}, "
|
||||
f"configured_horizon={self.config.action_horizon_flow}"
|
||||
)
|
||||
inputs_embeds = inputs_embeds.masked_scatter(action_mask, noisy_action_emb)
|
||||
|
||||
return inputs_embeds, flow, adarms_cond
|
||||
@@ -530,7 +544,10 @@ class ActionGenerationMixin(GenerationMixin):
|
||||
target_modules=target_modules,
|
||||
lora_dropout=lora_dropout,
|
||||
bias="none",
|
||||
task_type="CAUSAL_LM",
|
||||
# ``self.model`` is Wall-X's decoder, not the outer generation model.
|
||||
# Leave task_type unset so PEFT uses its generic wrapper and forwards
|
||||
# Wall-X's custom MoE arguments unchanged.
|
||||
task_type=None,
|
||||
)
|
||||
self.model = get_peft_model(self.model, config)
|
||||
# Log trainable parameter information
|
||||
|
||||
Reference in New Issue
Block a user