huangyanpei/VLA
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update Wall-X to 1.1.0 (#104)

Browse Source
This commit is contained in:
Starrick Liu
2026-06-15 11:40:00 +08:00
committed by GitHub
parent e23a586846
commit 72834e7de5
200 changed files with 33916 additions and 16771 deletions
Download Patch File Download Diff File Expand all files Collapse all files
wall_x/trainer/adapters/base_adapter.py
+494
View File
@@ -0,0 +1,494 @@
import logging
import os
from abc import ABC, abstractmethod
from collections import defaultdict
import torch
import torch.distributed as dist
from tqdm import tqdm
from wall_x.trainer.utils import move_batch_to_device
logger = logging.getLogger(__name__)
def load_trainer_data_config(cfg):
"""Return optional backend-specific trainer data config.
Public backends build datasets directly from TrainConfig. A backend that
needs an additional trainer data config may expose
``load_trainer_data_config(cfg)``.
"""
from wall_x.data import data_backend
backend = data_backend()
if backend.supports("load_trainer_data_config"):
return backend.load_trainer_data_config(cfg)
return None
class ModelAdapter(ABC):
"""Pure strategy that encapsulates model-type-specific behavior.
The adapter is stateless w.r.t. the trainer: every method declares
its inputs as parameters and communicates results via return values.
The trainer orchestrates calls and manages its own state.
Abstract methods (subclasses MUST implement):
load_processor -- load processor / tokenizer artifacts
build_model_config -- build model-specific config object
create_model -- instantiate the model
load_weights -- load pretrained weights & set normalizers
get_transformer_layer_cls -- layer classes for FSDP auto-wrap
load_dataset -- build dataset & training dataloader
forward -- run a forward pass
extract_loss -- extract scalar loss from model outputs
Concrete methods with default behaviour (override when needed):
load_step_and_epoch -- restore global_step / epoch from checkpoint
log_model_info -- log model-specific info (no-op)
collect_output_stats -- collect auxiliary loss / accuracy stats (no-op)
collect_param_norms -- compute per-group param L2 norms (no-op)
get_fwd_flops -- optional forward FLOPs hook
get_output_field -- retrieve a named field from outputs (static)
"""
def __init__(self, *, cfg=None, logger=None, model_type=None):
self.cfg = cfg
self.logger = logger
# ---- processor ----
@abstractmethod
def load_processor(self, action_statistic_dof):
"""Load processor and associated tokenizer artifacts.
Reads model/data config from ``self.cfg``.
Args:
action_statistic_dof: normalizer statistics dict.
Returns:
dict with keys:
"processor", "data_config", "tokenizer_mixin",
"train_action_tokenizer", "val_action_tokenizer",
"action_mapper", "num_added_tokens"
"""
...
# ---- model ----
@abstractmethod
def create_model(self, processor, tokenizer_mixin, model_config):
"""Instantiate the model (before weight loading / FSDP wrap).
Reads model config from ``self.cfg``.
"""
...
@abstractmethod
def build_model_config(self):
"""Build the model-specific config object from ``self.cfg.model``."""
...
@abstractmethod
def load_weights(self, model, normalizer_action, normalizer_propri, **kwargs):
"""Load pretrained weights, resize embeddings, set normalizers, etc."""
...
# ---- FSDP wrapping ----
@abstractmethod
def get_transformer_layer_cls(self):
"""Return layer classes for FSDP transformer_auto_wrap_policy."""
...
# ---- dataset ----
@abstractmethod
def load_dataset(self, data_config, processor, rank, world_size, **kwargs):
"""Load dataset and build the training dataloader.
Args:
data_config: parsed data config from load_processor().
processor: processor instance.
rank: current process rank.
world_size: total number of processes.
**kwargs: extra model-specific args.
Returns:
(dataset, train_dataloader, train_num) tuple.
"""
...
def load_step_and_epoch(self, checkpoint_path: str, is_incomplete_epoch: bool):
global_step = 0
start_epoch = 0
global_step_path = os.path.join(checkpoint_path, "global_step.pth")
if os.path.exists(global_step_path):
global_step = torch.load(global_step_path)["global_step"]
current_epoch_path = os.path.join(checkpoint_path, "current_epoch.pth")
if os.path.exists(current_epoch_path):
start_epoch = torch.load(current_epoch_path)["current_epoch"]
if not is_incomplete_epoch:
start_epoch = start_epoch + 1
return {"global_step": global_step, "start_epoch": start_epoch}
# ---- forward / loss ----
@abstractmethod
def forward(self, model, batch, **kwargs):
"""Run a forward pass.
Args:
model: the (possibly wrapped) model.
batch: dict of tensors already on device.
**kwargs: extra context (e.g. global_step, mode).
Returns:
Raw model outputs (dict or object).
"""
...
@abstractmethod
def extract_loss(self, outputs):
"""Extract the scalar training loss from model outputs.
Returns:
torch.Tensor (scalar).
"""
...
# ---- optional hooks (override when needed) ----
def log_model_info(self, model, log_fn):
"""Log model-specific info (e.g. attention implementation).
Args:
model: the unwrapped model.
log_fn: callable(str) for logging.
"""
pass
def collect_output_stats(
self,
outputs,
step_stats,
reduce_tensor_fn,
true_gather_fn,
tokenizer_mixin=None,
):
"""Collect model-output-specific statistics into step_stats.
Extracts auxiliary losses (cross_entropy_loss, flow_loss),
per-dataset channel losses, and action accuracy metrics from
the model outputs and merges them into *step_stats* in-place.
This is an optional hook - subclasses override it when their
output format carries extra fields. The default implementation
is a no-op.
Args:
outputs: raw model outputs (dict or object).
step_stats: dict being built by the trainer; mutated in-place.
reduce_tensor_fn: callable(tensor, average=True) -> scalar tensor,
wraps dist.all_reduce for the current parallelism.
true_gather_fn: callable(tensor) -> scalar | None,
gathers a value to rank-0 (returns None on
non-main ranks or when input is None).
tokenizer_mixin: optional ActionTokenizerMixin used to retrieve
extra accuracy keys (e.g. RVQ layer accuracy).
"""
pass
def collect_param_norms(
self, model, step_stats, device, reduce_tensor_fn, params_sharded=False
):
"""Compute per-group parameter L2 norms and merge into step_stats.
Groups parameters by name keyword:
- "visual" -> visual_param_norm
- "action" -> action_expert_param_norm
- (others) -> org_vlm_param_norm
- all -> total_param_norm
When parameters are sharded across ranks (FSDP FULL_SHARD /
HYBRID_SHARD), the local squared-sum is partial and must be
all_reduce'd (sum, not avg) before taking the square root.
When parameters are replicated (DDP, or FSDP with SHARD_GRAD_OP /
NO_SHARD), all_reduce would inflate the result by world_size.
This is an optional hook - the default implementation is a no-op.
Args:
model: the (possibly FSDP/DDP-wrapped) model.
step_stats: dict being built by the trainer; mutated in-place.
device: torch.device for tensor allocation.
reduce_tensor_fn: callable(tensor, average=False) -> tensor,
wraps dist.all_reduce for the current parallelism.
params_sharded: True when parameters are actually sharded across
ranks (FSDP FULL_SHARD / HYBRID_SHARD).
"""
pass
def collect_grad_norms(
self, model, step_stats, device, reduce_tensor_fn, params_sharded=False
):
"""Compute per-component gradient L2 norms and merge into step_stats.
Must be called while gradients are still available (before
optimizer.zero_grad). Under FSDP each rank holds a gradient shard,
so local squared-sums must be all_reduce'd (sum, not avg) before
taking the square root.
Default is a no-op; subclasses override for model-specific groupings.
"""
pass
def get_fwd_flops(self, autocast_context, model, batch, global_step):
"""Forward FLOPs profiling is disabled in the public package."""
return None
@staticmethod
def get_output_field(outputs, key, default=None):
"""Retrieve a named field from model outputs (dict or object)."""
if isinstance(outputs, dict):
return outputs.get(key, default)
return getattr(outputs, key, default)
# ---- console logging fields ----
def console_fields(self, tokenizer_mixin=None) -> list:
"""Per-step console fields as (stat_key, pretty_label, fmt_spec) triples.
MetricsLogger renders them in order, skipping any key not present
in step_stats. Public adapters may override this method to add
task-specific fields.
"""
fields = [
("video_loss", "vid_loss", ".6f"),
("action_loss", "act_loss", ".6f"),
("action_accuracy", "accuracy", ".4f"),
("flow_loss", "flow_loss", ".6f"),
]
if tokenizer_mixin is not None:
for key in tokenizer_mixin.get_accuracy_keys():
if key != "action_accuracy":
short_key = key.replace("action_accuracy_", "acc_")
fields.append((key, short_key, ".4f"))
fields.extend(
[
("action_grad_norm", "act_gnorm", ".4f"),
("video_grad_norm", "vid_gnorm", ".4f"),
]
)
return fields
# ---- distribution hooks ----
def convert_to_mix_precision_hint(
self,
model,
*,
device,
use_fsdp: bool,
log_fn=None,
):
"""Prepare *model* dtype + placement before FSDP/DDP wrap.
FSDP path: cast to fp32 on device; MixedPrecision policy handles the
bf16 compute cast at forward time. DDP path: delegate to the model's
``convert_to_mix_precision`` method (manual bf16 conversion, since
DDP has no equivalent of MixedPrecision), then move to device.
Adapters that handle placement themselves may override this hook.
"""
_log = log_fn or (lambda _msg: None)
if use_fsdp:
model.to(dtype=torch.float32)
_log("FSDP mode: all params fp32, MixedPrecision handles bf16 compute")
model.to(device)
else:
model.convert_to_mix_precision()
_log("DDP mode: params converted to bf16 (no MixedPrecision policy)")
model.to(device)
# ---- validation ----
def run_validation(
self,
*,
model,
val_dataloader,
rank,
world_size,
device,
autocast_context,
reduce_fn,
gather_fn,
logger,
global_step,
output_path,
tokenizer_mixin=None,
):
"""Default validation: forward-pass loss + adapter.collect_output_stats.
Public VLA adapters may override this for task-specific metrics.
``output_path`` is accepted for interface uniformity but not used
in this default.
"""
del output_path
model.eval()
log_dict = defaultdict(float)
pbar = tqdm(
val_dataloader,
desc="Validating",
total=len(val_dataloader),
disable=rank != 0,
)
for batch in pbar:
batch = move_batch_to_device(batch, device)
with autocast_context():
outputs = self.forward(model, batch, global_step=0)
loss = outputs["total_loss"] if "total_loss" in outputs else outputs["loss"]
log_dict["val_loss"] += reduce_fn(loss.detach()).item()
val_ce_loss = self.get_output_field(outputs, "cross_entropy_loss")
if val_ce_loss is not None:
log_ce_loss = gather_fn(val_ce_loss)
if log_ce_loss is not None and rank == 0:
log_dict["val_cross_entropy_loss"] += log_ce_loss
val_flow_loss = self.get_output_field(outputs, "flow_loss")
if val_flow_loss is not None:
log_flow_loss = gather_fn(val_flow_loss)
if log_flow_loss is not None and rank == 0:
log_dict["val_flow_loss"] += log_flow_loss
for k, v in collect_channel_loss_stats(
outputs,
prefix="val_",
tokenizer_mixin=tokenizer_mixin,
).items():
log_dict[k] += v
log_dict = {k: v / len(val_dataloader) for k, v in log_dict.items()}
if rank == 0:
if self.logger is not None:
self.logger.info(f"[FSDP Val] Step {global_step}: {log_dict}")
if logger is not None and hasattr(logger, "log"):
logger.log(log_dict, step=global_step)
val_info = f"[Validation] step {global_step}"
if "val_loss" in log_dict:
val_info += f" | val_loss {log_dict['val_loss']:.4f}"
if "val_cross_entropy_loss" in log_dict:
val_info += f" | val_ce_loss {log_dict['val_cross_entropy_loss']:.4f}"
if "val_flow_loss" in log_dict:
val_info += f" | val_flow_loss {log_dict['val_flow_loss']:.6f}"
if "val_action_accuracy" in log_dict:
val_info += f" | val_accuracy {log_dict['val_action_accuracy']:.4f}"
if tokenizer_mixin is not None:
for key in tokenizer_mixin.get_accuracy_keys():
if key != "action_accuracy":
val_key = f"val_{key}"
if val_key in log_dict:
short_key = key.replace("action_accuracy_", "acc_")
val_info += f" | val_{short_key} {log_dict[val_key]:.4f}"
if self.logger is not None:
self.logger.info(val_info)
model.train()
return log_dict
def init_validation(self, normalizer_action, normalizer_propri):
"""Optional pre-training-loop setup for validation. Default is a no-op."""
del normalizer_action, normalizer_propri
# ---- prediction / inference ----
def predict(
self,
prediction_type: str,
*,
model,
val_dataloader,
rank,
world_size,
device,
processor,
tokenizer_mixin,
logger,
current_step,
max_iteration=None,
save_dir=None,
max_samples=None,
):
"""Run an inference loop for the given ``prediction_type``.
Subclasses implement the dispatch table (flow_action / ar_action /
dllm_action / text / ...). Default raises - only VLAdapter has
meaningful predict paths today.
"""
raise NotImplementedError(
f"{type(self).__name__} does not support predict(prediction_type=...)"
)
# ---- optimizer-related defaults ----
@property
def default_action_lr_keywords(self) -> list[str]:
"""Keyword list used to identify action-expert parameters.
The trainer splits parameters into two groups when
``action_expert_learning_rate`` is set: names containing any of these
keywords go into the action-LR group. Subclasses override this when
the architecture uses different naming conventions.
"""
return [
"action_preprocessor",
"moe.experts.1",
"qkv_proj_experts.1",
"o_proj_experts.1",
"input_layernorms.1",
"post_attention_layernorms.1",
"model.norms.1",
]
def collect_channel_loss_stats(outputs, prefix="", tokenizer_mixin=None):
"""Gather per-dataset channel losses + accuracy metrics via dist.all_reduce.
Mirrors the legacy fsdp_trainer._collect_channel_loss_stats so training
and validation stats come out with identical keys. Returns {} when the
model output dict lacks ``channel_loss_dict``.
"""
stats = {}
channel_loss_dict = (
outputs.get("channel_loss_dict")
if isinstance(outputs, dict)
else getattr(outputs, "channel_loss_dict", None)
)
if channel_loss_dict is None:
return stats
channel_loss_count_dict = (
outputs.get("channel_loss_count_dict")
if isinstance(outputs, dict)
else getattr(outputs, "channel_loss_count_dict", None)
)
for dataset_name_i in channel_loss_dict:
count_tensor = channel_loss_count_dict[dataset_name_i].clone()
loss_tensor = channel_loss_dict[dataset_name_i].detach().clone()
dist.all_reduce(count_tensor, op=dist.ReduceOp.SUM)
dist.all_reduce(loss_tensor, op=dist.ReduceOp.SUM)
count_sum = count_tensor.item()
if count_sum >= 0.5:
stats[f"{prefix}channel_loss_{dataset_name_i}"] = (
loss_tensor.item() / count_sum
)
if "action_accuracy" in channel_loss_dict and tokenizer_mixin is not None:
acc_tensor = channel_loss_dict["action_accuracy"].detach().clone()
dist.all_reduce(acc_tensor, op=dist.ReduceOp.SUM)
world_size = dist.get_world_size() if dist.is_initialized() else 1
stats[f"{prefix}action_accuracy"] = acc_tensor.item() / world_size
for key in tokenizer_mixin.get_accuracy_keys():
if key != "action_accuracy" and key in channel_loss_dict:
rvq_acc_tensor = channel_loss_dict[key].detach().clone()
dist.all_reduce(rvq_acc_tensor, op=dist.ReduceOp.SUM)
stats[f"{prefix}{key}"] = rvq_acc_tensor.item() / world_size
return stats