Update Wall-X to 1.1.0 (#104)
This commit is contained in:
Starrick Liu
@@ -0,0 +1,128 @@
|
||||
"""Public model adapter registry."""
|
||||
|
||||
from importlib import import_module
|
||||
|
||||
from wall_x.trainer.adapters.base_adapter import ModelAdapter
|
||||
from wall_x.trainer.adapters.vla_model_adapter import VLAdapter
|
||||
|
||||
_ADAPTER_SPECS = {
|
||||
"qwen2_5": ("wall_x.model.qact.qwen2_5.adapter", "Qwen2_5Adapter"),
|
||||
}
|
||||
_ADAPTER_CLASS_EXPORTS = {
|
||||
"Qwen2_5Adapter": "qwen2_5",
|
||||
}
|
||||
ADAPTER_IMPORT_ERRORS: dict[str, str] = {}
|
||||
_ADAPTER_LOADING: set[str] = set()
|
||||
|
||||
|
||||
class _AdapterRegistry(dict[str, type[ModelAdapter]]):
|
||||
def get(self, model_type: str, default=None):
|
||||
_load_public_adapter(model_type)
|
||||
return dict.get(self, model_type, default)
|
||||
|
||||
def __contains__(self, model_type: object) -> bool:
|
||||
if isinstance(model_type, str):
|
||||
_load_public_adapter(model_type)
|
||||
return dict.__contains__(self, model_type)
|
||||
|
||||
def __getitem__(self, model_type: str) -> type[ModelAdapter]:
|
||||
_load_public_adapter(model_type)
|
||||
return dict.__getitem__(self, model_type)
|
||||
|
||||
def __iter__(self):
|
||||
_load_all_public_adapters()
|
||||
return dict.__iter__(self)
|
||||
|
||||
def keys(self):
|
||||
_load_all_public_adapters()
|
||||
return dict.keys(self)
|
||||
|
||||
def items(self):
|
||||
_load_all_public_adapters()
|
||||
return dict.items(self)
|
||||
|
||||
def values(self):
|
||||
_load_all_public_adapters()
|
||||
return dict.values(self)
|
||||
|
||||
|
||||
ADAPTER_REGISTRY: _AdapterRegistry = _AdapterRegistry()
|
||||
|
||||
|
||||
def _register_adapter(model_type: str, adapter_cls: type[ModelAdapter]) -> None:
|
||||
existing = dict.get(ADAPTER_REGISTRY, model_type)
|
||||
if existing is not None and existing is not adapter_cls:
|
||||
raise ValueError(
|
||||
f"model_type {model_type!r} already registered to "
|
||||
f"{existing.__name__}, cannot re-register to {adapter_cls.__name__}"
|
||||
)
|
||||
declared_model_type = getattr(adapter_cls, "MODEL_TYPE", model_type)
|
||||
if declared_model_type and declared_model_type != model_type:
|
||||
raise ValueError(
|
||||
f"adapter {adapter_cls.__name__} declares MODEL_TYPE="
|
||||
f"{declared_model_type!r}, but is registered as {model_type!r}"
|
||||
)
|
||||
dict.__setitem__(ADAPTER_REGISTRY, model_type, adapter_cls)
|
||||
|
||||
|
||||
def _load_public_adapter(model_type: str) -> type[ModelAdapter] | None:
|
||||
if dict.__contains__(ADAPTER_REGISTRY, model_type):
|
||||
return dict.__getitem__(ADAPTER_REGISTRY, model_type)
|
||||
spec = _ADAPTER_SPECS.get(model_type)
|
||||
if spec is None or model_type in _ADAPTER_LOADING:
|
||||
return None
|
||||
module_name, class_name = spec
|
||||
_ADAPTER_LOADING.add(model_type)
|
||||
try:
|
||||
module = import_module(module_name)
|
||||
adapter_cls = getattr(module, class_name)
|
||||
except (ImportError, AttributeError) as exc:
|
||||
ADAPTER_IMPORT_ERRORS[model_type] = f"{type(exc).__name__}: {exc}"
|
||||
return None
|
||||
finally:
|
||||
_ADAPTER_LOADING.discard(model_type)
|
||||
_register_adapter(model_type, adapter_cls)
|
||||
return adapter_cls
|
||||
|
||||
|
||||
def _load_all_public_adapters() -> None:
|
||||
for model_type in _ADAPTER_SPECS:
|
||||
_load_public_adapter(model_type)
|
||||
|
||||
|
||||
def format_adapter_error(model_type: str) -> str:
|
||||
msg = (
|
||||
f"Unsupported model type: {model_type}. "
|
||||
f"Registered: {sorted(ADAPTER_REGISTRY)}"
|
||||
)
|
||||
if model_type in ADAPTER_IMPORT_ERRORS:
|
||||
msg += f". Import failed: {ADAPTER_IMPORT_ERRORS[model_type]}"
|
||||
return msg
|
||||
|
||||
|
||||
def resolve_adapter(model_type: str) -> type[ModelAdapter]:
|
||||
"""Return a registered adapter class, loading lazily when needed."""
|
||||
adapter_cls = ADAPTER_REGISTRY.get(model_type)
|
||||
if adapter_cls is None:
|
||||
raise ValueError(format_adapter_error(model_type))
|
||||
return adapter_cls
|
||||
|
||||
|
||||
def __getattr__(name: str):
|
||||
model_type = _ADAPTER_CLASS_EXPORTS.get(name)
|
||||
if model_type is not None:
|
||||
adapter_cls = _load_public_adapter(model_type)
|
||||
if adapter_cls is not None:
|
||||
return adapter_cls
|
||||
raise AttributeError(name)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"ModelAdapter",
|
||||
"VLAdapter",
|
||||
"Qwen2_5Adapter",
|
||||
"ADAPTER_REGISTRY",
|
||||
"ADAPTER_IMPORT_ERRORS",
|
||||
"format_adapter_error",
|
||||
"resolve_adapter",
|
||||
]
|
||||
@@ -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
|
||||
@@ -0,0 +1,990 @@
|
||||
import os
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from tqdm import tqdm
|
||||
|
||||
from wall_x.trainer.adapters.base_adapter import ModelAdapter
|
||||
from wall_x.trainer.optimizer.dmuon import is_dmuon_model
|
||||
from wall_x.trainer.trainer_utils import (
|
||||
compute_action_metrics,
|
||||
load_qwen_pretrain_weight,
|
||||
load_wallx_processors_from_cfg,
|
||||
save_text_results_to_file,
|
||||
)
|
||||
|
||||
|
||||
class VLAdapter(ModelAdapter):
|
||||
"""Base adapter for public Qwen2.5 VLA models.
|
||||
|
||||
Concrete subclasses define the model/config classes, FSDP wrap layers,
|
||||
and optional task-specific logging hooks.
|
||||
"""
|
||||
|
||||
#: Subclass must set this. Identifies the variant in ``ADAPTER_REGISTRY``.
|
||||
MODEL_TYPE: str = ""
|
||||
|
||||
def __init__(self, *, cfg=None, logger=None, model_type=None):
|
||||
# Concrete subclasses set MODEL_TYPE; tests may pass model_type explicitly
|
||||
# (e.g. to swap a stub class into a different slot).
|
||||
resolved = model_type or self.MODEL_TYPE
|
||||
super().__init__(cfg=cfg, logger=logger, model_type=resolved)
|
||||
if not resolved:
|
||||
raise ValueError(
|
||||
f"{type(self).__name__} requires a MODEL_TYPE class attribute "
|
||||
"or an explicit model_type kwarg; got neither."
|
||||
)
|
||||
self.model_type = resolved
|
||||
self._dmuon_module = None
|
||||
self._dmuon_import_checked = False
|
||||
self._dmuon_named_param_cache = None
|
||||
|
||||
def _get_dmuon_module(self):
|
||||
if not self._dmuon_import_checked:
|
||||
try:
|
||||
import dmuon
|
||||
except ImportError:
|
||||
dmuon = None
|
||||
self._dmuon_module = dmuon
|
||||
self._dmuon_import_checked = True
|
||||
return self._dmuon_module
|
||||
|
||||
def _get_named_dmuon_dedicated_params(self, model):
|
||||
cache = self._dmuon_named_param_cache
|
||||
model_id = id(model)
|
||||
if cache is not None and cache[0] == model_id:
|
||||
return cache[1]
|
||||
|
||||
dmuon = self._get_dmuon_module()
|
||||
if dmuon is None:
|
||||
named_params = []
|
||||
else:
|
||||
module_to_name = {
|
||||
id(module): module_name for module_name, module in model.named_modules()
|
||||
}
|
||||
named_params = []
|
||||
for dparam in dmuon.get_dedicated_params(model):
|
||||
prefix = module_to_name.get(id(dparam.module), "")
|
||||
name = f"{prefix}.{dparam.param_name}" if prefix else dparam.param_name
|
||||
named_params.append((name, dparam))
|
||||
|
||||
self._dmuon_named_param_cache = (model_id, named_params)
|
||||
return named_params
|
||||
|
||||
# ---- variant hooks (subclass overrides) ----
|
||||
@classmethod
|
||||
def model_class(cls):
|
||||
"""Return the training-side model class for this variant."""
|
||||
raise NotImplementedError(f"{cls.__name__} must override model_class()")
|
||||
|
||||
@classmethod
|
||||
def config_class(cls):
|
||||
"""Return the HF PretrainedConfig class for this variant."""
|
||||
raise NotImplementedError(f"{cls.__name__} must override config_class()")
|
||||
|
||||
@classmethod
|
||||
def inference_model_class(cls):
|
||||
"""Return the model class to use for inference."""
|
||||
return cls.model_class()
|
||||
|
||||
# ---- helpers ----
|
||||
def _get_model_and_config_class(self):
|
||||
"""Resolve (ModelClass, ConfigClass) for this variant via classmethods."""
|
||||
cls = type(self)
|
||||
return cls.model_class(), cls.config_class()
|
||||
|
||||
# ---- processor ----
|
||||
def _build_processor_dict(self) -> dict:
|
||||
"""Flat-dict shape expected by legacy ``load_wallx_processors`` /
|
||||
``update_model_config`` / prediction loops.
|
||||
|
||||
Derived from typed TrainConfig. Kept as the one-place adapter layer
|
||||
between typed configs and legacy dict-based APIs; drop me when those
|
||||
downstream APIs are typed-ified.
|
||||
"""
|
||||
import dataclasses
|
||||
|
||||
flat = dataclasses.asdict(self.cfg.model)
|
||||
flat["model_type"] = self.cfg.model_type
|
||||
flat["data"] = dict(self.cfg._raw_data or {})
|
||||
flat["dof_config"] = self.cfg.task.dof_config
|
||||
flat["agent_pos_config"] = self.cfg.task.agent_pos_config
|
||||
if self.cfg.task.ar_dof_config is not None:
|
||||
flat["ar_dof_config"] = self.cfg.task.ar_dof_config
|
||||
flat["batch_size_per_gpu"] = self.cfg.hyperparams.batch_size_per_gpu
|
||||
return flat
|
||||
|
||||
def load_processor(self, action_statistic_dof):
|
||||
processors_dict = load_wallx_processors_from_cfg(
|
||||
self.cfg,
|
||||
normalizer=getattr(self, "normalizer_action", None),
|
||||
action_statistic_dof=action_statistic_dof,
|
||||
)
|
||||
self.logger.info(
|
||||
f"processor vocab size: {len(processors_dict['processor'].tokenizer.vocab)}"
|
||||
)
|
||||
self.logger.info(
|
||||
f"num added tokens to processor: {processors_dict['num_added_tokens']}"
|
||||
)
|
||||
return {
|
||||
"processor": processors_dict["processor"],
|
||||
"data_config": self._build_processor_dict(),
|
||||
"tokenizer_mixin": processors_dict.get("tokenizer_mixin"),
|
||||
"train_action_tokenizer": processors_dict["train_action_tokenizer"],
|
||||
"val_action_tokenizer": processors_dict["val_action_tokenizer"],
|
||||
"action_mapper": processors_dict["action_mapper"],
|
||||
"num_added_tokens": processors_dict["num_added_tokens"],
|
||||
}
|
||||
|
||||
# ---- model ----
|
||||
def build_model_config(self):
|
||||
_, ConfigClass = self._get_model_and_config_class()
|
||||
qwen_vl_act_config_path = self.cfg.model.config_path
|
||||
if qwen_vl_act_config_path.endswith(".json"):
|
||||
model_config = ConfigClass.from_json_file(qwen_vl_act_config_path)
|
||||
else:
|
||||
model_config = ConfigClass.from_pretrained(qwen_vl_act_config_path)
|
||||
|
||||
assert self.model_type in model_config.model_type, (
|
||||
f"Mismatch of model type: model type in config file is "
|
||||
f"{model_config.model_type}, but the model type is {self.model_type}."
|
||||
)
|
||||
|
||||
model_config.update_model_config(self._build_processor_dict())
|
||||
return model_config
|
||||
|
||||
def create_model(self, processor, tokenizer_mixin, model_config):
|
||||
ModelClass, _ = self._get_model_and_config_class()
|
||||
use_selective_recompute = self.cfg.distributed.use_selective_recompute
|
||||
return ModelClass(
|
||||
model_config,
|
||||
processor,
|
||||
tokenizer_mixin=tokenizer_mixin,
|
||||
use_selective_recompute=use_selective_recompute,
|
||||
)
|
||||
|
||||
def load_weights(self, model, normalizer_action, normalizer_propri, **kwargs):
|
||||
import copy
|
||||
|
||||
processor = kwargs.get("processor")
|
||||
|
||||
if self.cfg.model.pretrained_path:
|
||||
model, err = load_qwen_pretrain_weight(
|
||||
model, self.cfg.model.pretrained_path
|
||||
)
|
||||
|
||||
if processor is not None:
|
||||
model.resize_token_embeddings(len(processor.tokenizer))
|
||||
|
||||
if hasattr(model, "enable_input_require_grads"):
|
||||
model.enable_input_require_grads()
|
||||
elif (
|
||||
hasattr(model, "get_input_embeddings")
|
||||
and model.get_input_embeddings() is not None
|
||||
):
|
||||
|
||||
def _make_inputs_require_grad(module, input, output):
|
||||
output.requires_grad_(True)
|
||||
|
||||
model.get_input_embeddings().register_forward_hook(
|
||||
_make_inputs_require_grad
|
||||
)
|
||||
|
||||
if hasattr(model, "set_normalizer"):
|
||||
model.set_normalizer(
|
||||
copy.deepcopy(normalizer_action),
|
||||
copy.deepcopy(normalizer_propri),
|
||||
)
|
||||
|
||||
return model
|
||||
|
||||
# ---- FSDP wrapping ----
|
||||
def get_transformer_layer_cls(self):
|
||||
"""Return the FSDP transformer-wrap layer classes for this variant.
|
||||
|
||||
Subclass override - see the per-variant adapter for the actual
|
||||
layer classes.
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
f"{type(self).__name__} must override get_transformer_layer_cls()"
|
||||
)
|
||||
|
||||
# ---- dataset ----
|
||||
def load_dataset(self, data_config, processor, rank, world_size, **kwargs):
|
||||
"""Dispatch to the backend named by ``cfg.data.dataset_type``.
|
||||
|
||||
The backend-specific wiring (resume indices, pool offsets,
|
||||
processor wrapping) lives inside each backend's ``build()``;
|
||||
this method only assembles a :class:`BuildContext` and forwards.
|
||||
"""
|
||||
import copy
|
||||
|
||||
from wall_x.data import BuildContext, build_data
|
||||
|
||||
resume_state = None
|
||||
resume_batches = kwargs.get("resume_batches", 0)
|
||||
indices = kwargs.get("resume_indices")
|
||||
if indices is None and self.cfg.checkpoint.resume_from:
|
||||
checkpoint_path = self.cfg.checkpoint.resume_from
|
||||
if os.path.isfile(checkpoint_path):
|
||||
checkpoint_path = os.path.dirname(checkpoint_path)
|
||||
if os.path.isdir(checkpoint_path):
|
||||
indices = self._load_episode_indices(checkpoint_path, rank)
|
||||
|
||||
if resume_batches or indices or kwargs.get("episode_container_checkpoint_path"):
|
||||
resume_state = {
|
||||
"resume_batches": resume_batches,
|
||||
"indices": indices,
|
||||
"episode_container_checkpoint_path": kwargs.get(
|
||||
"episode_container_checkpoint_path"
|
||||
),
|
||||
}
|
||||
|
||||
ctx = BuildContext(
|
||||
rank=rank,
|
||||
world_size=world_size,
|
||||
processor=processor,
|
||||
tokenizer_mixin=kwargs.get("tokenizer_mixin"),
|
||||
normalizer_action=copy.deepcopy(kwargs.get("normalizer_action")),
|
||||
normalizer_propri=copy.deepcopy(kwargs.get("normalizer_propri")),
|
||||
model_config=kwargs.get("model_config"),
|
||||
resume_state=resume_state,
|
||||
)
|
||||
|
||||
bundle = build_data(self.cfg, ctx)
|
||||
return bundle.dataset, bundle.train_loader, bundle.train_iters
|
||||
|
||||
# ---- forward / loss ----
|
||||
def forward(self, model, batch, **kwargs):
|
||||
self._last_batch_info = self._extract_batch_info(batch)
|
||||
mode = kwargs.get("mode", "train")
|
||||
return model(**batch, mode=mode)
|
||||
|
||||
def extract_loss(self, outputs):
|
||||
return outputs.loss
|
||||
|
||||
def collect_output_stats(
|
||||
self,
|
||||
outputs,
|
||||
step_stats,
|
||||
reduce_tensor_fn,
|
||||
true_gather_fn,
|
||||
tokenizer_mixin=None,
|
||||
):
|
||||
"""Collect VL-specific auxiliary losses and accuracy metrics.
|
||||
|
||||
Handles:
|
||||
- cross_entropy_loss / flow_loss (all-reduce + true_gather)
|
||||
- per-dataset channel losses (all-reduce via count-weighted avg)
|
||||
- action_accuracy and extra RVQ layer accuracies
|
||||
"""
|
||||
import torch.distributed as dist
|
||||
|
||||
cross_entropy_loss = self.get_output_field(outputs, "cross_entropy_loss")
|
||||
if cross_entropy_loss is not None:
|
||||
step_stats["cross_entropy_loss"] = reduce_tensor_fn(
|
||||
cross_entropy_loss
|
||||
).item()
|
||||
|
||||
flow_loss = self.get_output_field(outputs, "flow_loss")
|
||||
if flow_loss is not None:
|
||||
step_stats["flow_loss"] = reduce_tensor_fn(flow_loss).item()
|
||||
|
||||
log_ce_loss = true_gather_fn(
|
||||
self.get_output_field(outputs, "cross_entropy_loss")
|
||||
)
|
||||
if log_ce_loss is not None:
|
||||
step_stats["cross_entropy_loss"] = log_ce_loss
|
||||
|
||||
log_flow_loss = true_gather_fn(self.get_output_field(outputs, "flow_loss"))
|
||||
if log_flow_loss is not None:
|
||||
step_stats["flow_loss"] = log_flow_loss
|
||||
|
||||
channel_loss_dict = self.get_output_field(outputs, "channel_loss_dict")
|
||||
channel_loss_count_dict = self.get_output_field(
|
||||
outputs, "channel_loss_count_dict"
|
||||
)
|
||||
if channel_loss_dict is not 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)
|
||||
|
||||
cout_sum = count_tensor.item()
|
||||
if cout_sum >= 0.5:
|
||||
step_stats[f"channel_loss_{dataset_name_i}"] = (
|
||||
loss_tensor.item() / cout_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
|
||||
step_stats["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)
|
||||
step_stats[key] = rvq_acc_tensor.item() / world_size
|
||||
|
||||
def collect_param_norms(
|
||||
self, model, step_stats, device, reduce_tensor_fn, params_sharded=False
|
||||
):
|
||||
"""Compute per-group parameter L2 norms for VL models.
|
||||
|
||||
Groups:
|
||||
"visual" -> visual_param_norm
|
||||
"action" -> action_expert_param_norm
|
||||
others -> org_vlm_param_norm
|
||||
all -> total_param_norm
|
||||
"""
|
||||
total_sq = torch.tensor(0.0, dtype=torch.float32, device=device)
|
||||
visual_sq = torch.tensor(0.0, dtype=torch.float32, device=device)
|
||||
action_sq = torch.tensor(0.0, dtype=torch.float32, device=device)
|
||||
org_vlm_sq = torch.tensor(0.0, dtype=torch.float32, device=device)
|
||||
dedicated_sq = torch.tensor(0.0, dtype=torch.float32, device=device)
|
||||
dedicated_params = []
|
||||
|
||||
def add_named_sq(name, sq):
|
||||
nonlocal total_sq, visual_sq, action_sq, org_vlm_sq
|
||||
total_sq += sq
|
||||
if "visual" in name:
|
||||
visual_sq += sq
|
||||
elif "action" in name:
|
||||
action_sq += sq
|
||||
else:
|
||||
org_vlm_sq += sq
|
||||
|
||||
with torch.no_grad():
|
||||
for name, param in model.named_parameters():
|
||||
if not param.requires_grad:
|
||||
continue
|
||||
if param.numel() == 0 or hasattr(param, "_dedicated_owner_rank"):
|
||||
continue
|
||||
tensor = param.detach()
|
||||
if hasattr(tensor, "to_local"):
|
||||
tensor = tensor.to_local()
|
||||
if tensor.numel() == 0:
|
||||
continue
|
||||
sq = torch.sum(tensor.float() ** 2)
|
||||
add_named_sq(name, sq)
|
||||
|
||||
if is_dmuon_model(model):
|
||||
dedicated_params = self._get_named_dmuon_dedicated_params(model)
|
||||
# DMuon replaces dedicated params with placeholders in
|
||||
# named_parameters(). Count the authoritative dedicated
|
||||
# storage here so total/category norms cover the full model.
|
||||
for name, dparam in dedicated_params:
|
||||
if getattr(dparam, "_dmuon_route", None) == "sharded_adamw":
|
||||
replicate_group = getattr(dparam, "replicate_group", None)
|
||||
# replicate_group=None is 1D shard-only mode: each rank
|
||||
# contributes a distinct shard and the later reduce sums
|
||||
# them into the global norm.
|
||||
if (
|
||||
replicate_group is not None
|
||||
and replicate_group.rank()
|
||||
!= getattr(dparam, "owner_replicate", 0)
|
||||
):
|
||||
continue
|
||||
tensor = getattr(dparam, "_sharded_adamw_data", None)
|
||||
if tensor is None:
|
||||
continue
|
||||
valid_numel = int(
|
||||
getattr(
|
||||
dparam,
|
||||
"_sharded_adamw_valid_numel",
|
||||
tensor.numel(),
|
||||
)
|
||||
)
|
||||
tensor = tensor[:valid_numel]
|
||||
else:
|
||||
if not bool(getattr(dparam, "is_owner", False)):
|
||||
continue
|
||||
tensor = getattr(dparam, "_owned_data", None)
|
||||
if tensor is None:
|
||||
continue
|
||||
|
||||
if tensor.numel() == 0:
|
||||
continue
|
||||
sq = torch.sum(tensor.detach().float() ** 2)
|
||||
dedicated_sq += sq
|
||||
add_named_sq(name, sq)
|
||||
|
||||
if params_sharded:
|
||||
total_sq = reduce_tensor_fn(total_sq, average=False)
|
||||
visual_sq = reduce_tensor_fn(visual_sq, average=False)
|
||||
action_sq = reduce_tensor_fn(action_sq, average=False)
|
||||
org_vlm_sq = reduce_tensor_fn(org_vlm_sq, average=False)
|
||||
dedicated_sq = reduce_tensor_fn(dedicated_sq, average=False)
|
||||
|
||||
step_stats["total_param_norm"] = torch.sqrt(total_sq).item()
|
||||
step_stats["visual_param_norm"] = torch.sqrt(visual_sq).item()
|
||||
step_stats["action_expert_param_norm"] = torch.sqrt(action_sq).item()
|
||||
step_stats["org_vlm_param_norm"] = torch.sqrt(org_vlm_sq).item()
|
||||
if dedicated_params:
|
||||
step_stats["dmuon_dedicated_param_norm"] = torch.sqrt(dedicated_sq).item()
|
||||
|
||||
# ---- MFU computation ----
|
||||
@staticmethod
|
||||
def _extract_batch_info(batch):
|
||||
"""Extract token counts from a training batch."""
|
||||
info = {}
|
||||
input_ids = batch.get("input_ids")
|
||||
if input_ids is not None:
|
||||
info["batch_size"] = input_ids.shape[0]
|
||||
info["seq_length"] = input_ids.shape[1]
|
||||
|
||||
moe_token_types = batch.get("moe_token_types")
|
||||
if moe_token_types is not None:
|
||||
info["num_lang_tokens"] = (
|
||||
int((moe_token_types == 0).sum().item()) // info["batch_size"]
|
||||
)
|
||||
info["num_action_tokens"] = (
|
||||
int((moe_token_types == 1).sum().item()) // info["batch_size"]
|
||||
)
|
||||
else:
|
||||
info["num_lang_tokens"] = info.get("seq_length", 0)
|
||||
info["num_action_tokens"] = 0
|
||||
|
||||
pixel_values = batch.get("pixel_values")
|
||||
info["vision_seq_length"] = (
|
||||
pixel_values.shape[0] if pixel_values is not None else 0
|
||||
)
|
||||
|
||||
labels = batch.get("labels")
|
||||
if labels is not None:
|
||||
info["num_loss_tokens"] = int((labels[..., 1:] != -100).sum().item())
|
||||
else:
|
||||
info["num_loss_tokens"] = None
|
||||
|
||||
return info
|
||||
|
||||
def _compute_detailed_flops(self, model, config):
|
||||
"""Compute detailed per-module FLOPs for one training step.
|
||||
|
||||
Uses the same formulas as scripts/profile_forward.py:compute_module_flops,
|
||||
with fwd+bwd multipliers:
|
||||
- Frozen modules (e.g. ViT if frozen): 1x forward
|
||||
- Trainable modules: 3x forward (fwd + 2x bwd)
|
||||
"""
|
||||
|
||||
batch_info = getattr(self, "_last_batch_info", None)
|
||||
if not batch_info:
|
||||
return None
|
||||
|
||||
model_config = model.config if hasattr(model, "config") else None
|
||||
if model_config is None:
|
||||
return None
|
||||
|
||||
B = batch_info["batch_size"]
|
||||
S = batch_info["seq_length"]
|
||||
num_lang = batch_info["num_lang_tokens"]
|
||||
num_act = batch_info["num_action_tokens"]
|
||||
N_lang = B * num_lang
|
||||
N_act = B * num_act
|
||||
N_total = N_lang + N_act
|
||||
Nv = batch_info["vision_seq_length"]
|
||||
num_loss_tokens = batch_info["num_loss_tokens"]
|
||||
|
||||
H = model_config.hidden_size
|
||||
num_layers = model_config.num_hidden_layers
|
||||
num_heads = model_config.num_attention_heads
|
||||
num_kv = model_config.num_key_value_heads
|
||||
vocab_size = getattr(model_config, "padded_vocab_size", model_config.vocab_size)
|
||||
|
||||
use_mot = getattr(model_config, "attention_moe", False)
|
||||
use_moe_mlp = getattr(model_config, "mlp_moe", False)
|
||||
dim_inputs = getattr(model_config, "dim_inputs", (H, H))
|
||||
dim_lang, dim_act = dim_inputs
|
||||
|
||||
grad_accum = self.cfg.hyperparams.gradient_accumulation_steps
|
||||
|
||||
vit_frozen = True
|
||||
if hasattr(model, "visual"):
|
||||
for p in model.visual.parameters():
|
||||
if p.requires_grad:
|
||||
vit_frozen = False
|
||||
break
|
||||
vit_mult = 1.0 if vit_frozen else 3.0
|
||||
|
||||
vit_fwd_flops = 0
|
||||
if hasattr(model_config, "vision_config") and Nv > 0:
|
||||
vcfg = model_config.vision_config
|
||||
Hv = vcfg.hidden_size
|
||||
Iv = vcfg.intermediate_size
|
||||
out_hidden = vcfg.out_hidden_size
|
||||
depth_v = vcfg.depth
|
||||
|
||||
F_qkv_v = 6 * Nv * Hv * Hv
|
||||
F_o_v = 2 * Nv * Hv * Hv
|
||||
F_mlp_v = 6 * Nv * Hv * Iv + 2 * Nv * Iv
|
||||
F_linear_per_layer = F_qkv_v + F_o_v + F_mlp_v
|
||||
|
||||
fullatt_set = set(getattr(vcfg, "fullatt_block_indexes", []))
|
||||
num_images = max(Nv // 768, 1)
|
||||
si = Nv // num_images
|
||||
sum_si_sq = num_images * si * si
|
||||
sum_wi_sq = sum_si_sq // 16
|
||||
|
||||
for i in range(depth_v):
|
||||
if i in fullatt_set:
|
||||
vit_fwd_flops += F_linear_per_layer + 4 * Hv * sum_si_sq
|
||||
else:
|
||||
vit_fwd_flops += F_linear_per_layer + 4 * Hv * sum_wi_sq
|
||||
|
||||
merge_unit = getattr(vcfg, "spatial_merge_size", 2) ** 2
|
||||
merger_hidden = Hv * merge_unit
|
||||
Nv_merged = Nv // merge_unit
|
||||
vit_fwd_flops += (
|
||||
2 * Nv_merged * merger_hidden * merger_hidden
|
||||
+ 2 * Nv_merged * merger_hidden * out_hidden
|
||||
)
|
||||
|
||||
kv_ratio = num_kv / num_heads
|
||||
F_matmul_fwd = 4 * B * (S**2) * H
|
||||
|
||||
if not use_mot:
|
||||
F_attn_fwd = (
|
||||
2 * N_total * H * H
|
||||
+ 4 * N_total * H * H * kv_ratio
|
||||
+ 2 * N_total * H * H
|
||||
+ F_matmul_fwd
|
||||
)
|
||||
else:
|
||||
F_attn_fwd = (
|
||||
N_lang * dim_lang * H * (4 + 4 * kv_ratio)
|
||||
+ N_act * dim_act * H * (4 + 4 * kv_ratio)
|
||||
+ F_matmul_fwd
|
||||
)
|
||||
|
||||
if not use_moe_mlp:
|
||||
ffn_hidden = model_config.intermediate_size
|
||||
F_mlp_fwd = 6 * N_total * H * ffn_hidden + 2 * N_total * ffn_hidden
|
||||
else:
|
||||
hid_lang = model_config.experts[0]["intermediate_size"]
|
||||
hid_act = model_config.experts[1]["intermediate_size"]
|
||||
F_mlp_fwd = (6 * N_lang * dim_lang * hid_lang + 2 * N_lang * hid_lang) + (
|
||||
6 * N_act * dim_act * hid_act + 2 * N_act * hid_act
|
||||
)
|
||||
|
||||
decoder_fwd_flops = num_layers * (F_attn_fwd + F_mlp_fwd)
|
||||
|
||||
N_lm = num_loss_tokens if num_loss_tokens is not None else N_total
|
||||
lm_head_fwd_flops = 2 * N_lm * H * vocab_size
|
||||
|
||||
total_flops = (
|
||||
vit_mult * vit_fwd_flops + 3.0 * decoder_fwd_flops + 3.0 * lm_head_fwd_flops
|
||||
) * grad_accum
|
||||
|
||||
return {
|
||||
"total_flops": total_flops,
|
||||
"vit_fwd_flops": vit_fwd_flops,
|
||||
"decoder_fwd_flops": decoder_fwd_flops,
|
||||
"lm_head_fwd_flops": lm_head_fwd_flops,
|
||||
"vit_mult": vit_mult,
|
||||
}
|
||||
|
||||
def compute_mfu(self, model, step_time_seconds):
|
||||
"""Compute detailed Model FLOPs Utilization for VL transformer training."""
|
||||
if step_time_seconds <= 0:
|
||||
return None
|
||||
try:
|
||||
config = self._build_processor_dict()
|
||||
flops_info = self._compute_detailed_flops(model, config)
|
||||
if flops_info is None:
|
||||
return None
|
||||
|
||||
gpu_peak_tflops = 312.0 # TODO: move to DebugConfig if needed
|
||||
peak_flops = gpu_peak_tflops * 1e12
|
||||
|
||||
# Data-parallel: each GPU independently processes its own micro-batch,
|
||||
# so per-GPU MFU = per_gpu_flops / (step_time * per_gpu_peak).
|
||||
# No division by num_gpus needed.
|
||||
total_flops = flops_info["total_flops"]
|
||||
mfu = total_flops / (step_time_seconds * peak_flops)
|
||||
|
||||
return {
|
||||
"mfu": mfu,
|
||||
"flops_per_step_T": total_flops / 1e12,
|
||||
}
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
# ---- optional hooks ----
|
||||
@staticmethod
|
||||
def log_attention_implementation(logger, model):
|
||||
"""Log the attention implementation name. Variant-specific because the
|
||||
model layout (model.model vs model.model.language_model) differs."""
|
||||
raise NotImplementedError(
|
||||
"VLAdapter subclasses must override log_attention_implementation()"
|
||||
)
|
||||
|
||||
def _load_episode_indices(self, checkpoint_path: str, rank: int):
|
||||
"""Read backend-managed per-rank resume offsets if supported."""
|
||||
from wall_x.data import data_backend
|
||||
|
||||
backend = data_backend()
|
||||
if not backend.supports("load_episode_indices"):
|
||||
return None
|
||||
return backend.load_episode_indices(checkpoint_path, rank)
|
||||
|
||||
# ---- 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,
|
||||
):
|
||||
"""Dispatch VLA inference by prediction_type.
|
||||
|
||||
flow_action / dllm_action return per-rank L1 action metrics; ar_action
|
||||
is reserved (was unimplemented pre-refactor too); text runs text
|
||||
generation with optional point-L1 distance scoring.
|
||||
"""
|
||||
config = self._build_processor_dict()
|
||||
if prediction_type == "flow_action":
|
||||
return self._predict_flow_action(
|
||||
model=model,
|
||||
val_dataloader=val_dataloader,
|
||||
rank=rank,
|
||||
world_size=world_size,
|
||||
device=device,
|
||||
config=config,
|
||||
logger=logger,
|
||||
current_step=current_step,
|
||||
max_iteration=max_iteration,
|
||||
)
|
||||
if prediction_type == "ar_action":
|
||||
return None # matches pre-refactor behaviour (was pass)
|
||||
if prediction_type == "dllm_action":
|
||||
return self._predict_dllm_action(
|
||||
model=model,
|
||||
val_dataloader=val_dataloader,
|
||||
rank=rank,
|
||||
world_size=world_size,
|
||||
device=device,
|
||||
config=config,
|
||||
processor=processor,
|
||||
tokenizer_mixin=tokenizer_mixin,
|
||||
logger=logger,
|
||||
current_step=current_step,
|
||||
max_iteration=max_iteration,
|
||||
)
|
||||
if prediction_type == "text":
|
||||
return self._predict_text(
|
||||
model=model,
|
||||
val_dataloader=val_dataloader,
|
||||
rank=rank,
|
||||
device=device,
|
||||
logger=logger,
|
||||
current_step=current_step,
|
||||
max_samples=max_samples,
|
||||
save_dir=save_dir,
|
||||
)
|
||||
raise ValueError(f"Unsupported prediction type: {prediction_type}")
|
||||
|
||||
@torch.no_grad()
|
||||
def _predict_flow_action(
|
||||
self,
|
||||
*,
|
||||
model,
|
||||
val_dataloader,
|
||||
rank,
|
||||
world_size,
|
||||
device,
|
||||
config,
|
||||
logger,
|
||||
current_step,
|
||||
max_iteration,
|
||||
):
|
||||
if dist.is_initialized():
|
||||
dist.barrier()
|
||||
|
||||
total_num = len(val_dataloader)
|
||||
if max_iteration:
|
||||
total_num = min(max_iteration, total_num)
|
||||
model.eval()
|
||||
|
||||
all_preds, all_actions = [], []
|
||||
pepoch = tqdm(
|
||||
total=total_num,
|
||||
desc=f"Predicting ckpt at step {current_step}",
|
||||
disable=rank != 0,
|
||||
)
|
||||
for batch_idx, batch in enumerate(val_dataloader):
|
||||
if batch_idx >= total_num:
|
||||
break
|
||||
batch = _move_batch(batch, device)
|
||||
|
||||
model_output = model.generate_flow_action(
|
||||
action_horizon=config["data"]["action_horizon_flow"],
|
||||
action_dim=model.action_preprocessor.action_dim,
|
||||
**batch,
|
||||
)
|
||||
pred_action, gt_action = (
|
||||
model_output["predict_action"],
|
||||
model_output["gt_action"],
|
||||
)
|
||||
|
||||
pred_list = [torch.zeros_like(pred_action) for _ in range(world_size)]
|
||||
gt_list = [torch.zeros_like(gt_action) for _ in range(world_size)]
|
||||
dist.all_gather(pred_list, pred_action)
|
||||
dist.all_gather(gt_list, gt_action)
|
||||
|
||||
if dist.is_initialized():
|
||||
dist.barrier()
|
||||
if rank == 0:
|
||||
all_preds.append(torch.cat(pred_list, dim=0).cpu())
|
||||
all_actions.append(torch.cat(gt_list, dim=0).cpu())
|
||||
pepoch.update(1)
|
||||
|
||||
pepoch.close()
|
||||
|
||||
if rank == 0:
|
||||
step_log = {}
|
||||
if all_preds:
|
||||
all_preds = torch.cat(all_preds, dim=0)
|
||||
all_actions = torch.cat(all_actions, dim=0)
|
||||
step_log = compute_action_metrics(
|
||||
all_preds, all_actions, config=config, step_log=step_log
|
||||
)
|
||||
if logger and step_log:
|
||||
logger.log(step_log, step=current_step)
|
||||
if self.logger is not None:
|
||||
self.logger.info(
|
||||
f"Step {current_step}, Validation L1 Loss: {step_log.get('val_action_l1', -1)}"
|
||||
)
|
||||
return step_log
|
||||
return None
|
||||
|
||||
@torch.no_grad()
|
||||
def _predict_dllm_action(
|
||||
self,
|
||||
*,
|
||||
model,
|
||||
val_dataloader,
|
||||
rank,
|
||||
world_size,
|
||||
device,
|
||||
config,
|
||||
processor,
|
||||
tokenizer_mixin,
|
||||
logger,
|
||||
current_step,
|
||||
max_iteration,
|
||||
):
|
||||
if dist.is_initialized():
|
||||
dist.barrier()
|
||||
|
||||
total_num = len(val_dataloader)
|
||||
if max_iteration:
|
||||
total_num = min(max_iteration, total_num)
|
||||
model.eval()
|
||||
|
||||
all_preds, all_actions = [], []
|
||||
pepoch = tqdm(
|
||||
total=total_num,
|
||||
desc=f"Predicting ckpt at step {current_step}",
|
||||
disable=rank != 0,
|
||||
)
|
||||
for batch_idx, batch in enumerate(val_dataloader):
|
||||
if batch_idx >= total_num:
|
||||
break
|
||||
batch = _move_batch(batch, device)
|
||||
batch = self._preprocess_dllm_batch(batch, processor, tokenizer_mixin)
|
||||
|
||||
if hasattr(model.action_tokenizer, "max_waypoints"):
|
||||
total_ar_step = model.action_tokenizer.max_waypoints
|
||||
elif hasattr(model.action_tokenizer, "max_length"):
|
||||
total_ar_step = model.action_tokenizer.max_length
|
||||
else:
|
||||
raise ValueError(
|
||||
"Unknown action_tokenizer type for dllm action inference"
|
||||
)
|
||||
model_output = model.generate_dllm_action(
|
||||
action_dim=7,
|
||||
action_horizon=config["data"]["action_horizon_flow"],
|
||||
use_ar_action=False,
|
||||
total_ar_step=total_ar_step,
|
||||
**batch,
|
||||
)
|
||||
pred_action, gt_action = (
|
||||
model_output["predict_action"],
|
||||
model_output["gt_action"],
|
||||
)
|
||||
|
||||
pred_list = [torch.zeros_like(pred_action) for _ in range(world_size)]
|
||||
gt_list = [torch.zeros_like(gt_action) for _ in range(world_size)]
|
||||
dist.all_gather(pred_list, pred_action)
|
||||
dist.all_gather(gt_list, gt_action)
|
||||
|
||||
if dist.is_initialized():
|
||||
dist.barrier()
|
||||
if rank == 0:
|
||||
all_preds.append(torch.cat(pred_list, dim=0).cpu())
|
||||
all_actions.append(torch.cat(gt_list, dim=0).cpu())
|
||||
pepoch.update(1)
|
||||
|
||||
pepoch.close()
|
||||
|
||||
if rank == 0:
|
||||
step_log = {}
|
||||
if all_preds:
|
||||
all_preds = torch.cat(all_preds, dim=0)
|
||||
all_actions = torch.cat(all_actions, dim=0)
|
||||
step_log = compute_action_metrics(
|
||||
all_preds, all_actions, config=config, step_log=step_log
|
||||
)
|
||||
if logger and step_log:
|
||||
logger.log(step_log, step=current_step)
|
||||
if self.logger is not None:
|
||||
self.logger.info(
|
||||
f"Step {current_step}, Validation L1 Loss: {step_log.get('val_action_l1', -1)}"
|
||||
)
|
||||
return step_log
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _preprocess_dllm_batch(batch, processor, tokenizer_mixin):
|
||||
"""Patch the dllm action placeholders into input_ids (in place)."""
|
||||
input_ids = batch["input_ids"]
|
||||
prefix_length = batch["prefix_length"]
|
||||
placeholder_str = tokenizer_mixin.get_placeholder_for_dllm()
|
||||
placeholder_seq = torch.tensor(
|
||||
processor.tokenizer.convert_tokens_to_ids(placeholder_str)
|
||||
)
|
||||
placeholder_len = placeholder_seq.shape[0]
|
||||
input_ids[:, prefix_length - placeholder_len - 2 : prefix_length - 2] = (
|
||||
placeholder_seq
|
||||
)
|
||||
batch.update({"input_ids": input_ids})
|
||||
return batch
|
||||
|
||||
@torch.no_grad()
|
||||
def _predict_text(
|
||||
self,
|
||||
*,
|
||||
model,
|
||||
val_dataloader,
|
||||
rank,
|
||||
device,
|
||||
logger,
|
||||
current_step,
|
||||
max_samples,
|
||||
save_dir,
|
||||
):
|
||||
from wall_x._vendor.x2robot_utils.grounding import calculate_point_l1_distance
|
||||
|
||||
if dist.is_initialized():
|
||||
dist.barrier()
|
||||
|
||||
total_num = (
|
||||
len(val_dataloader)
|
||||
if max_samples is None
|
||||
else min(max_samples, len(val_dataloader))
|
||||
)
|
||||
model.eval()
|
||||
|
||||
all_input_texts, all_gt_texts, all_pred_texts = [], [], []
|
||||
all_point_l1_distances = []
|
||||
|
||||
pepoch = tqdm(
|
||||
total=total_num,
|
||||
desc=f"Predicting text at step {current_step}",
|
||||
disable=rank != 0,
|
||||
)
|
||||
|
||||
for batch_idx, batch in enumerate(val_dataloader):
|
||||
if batch_idx >= total_num:
|
||||
break
|
||||
batch = _move_batch(batch, device)
|
||||
|
||||
model_output = model.generate_text(
|
||||
input_ids=batch.get("input_ids"),
|
||||
attention_mask=batch.get("attention_mask"),
|
||||
moe_token_types=batch.get("moe_token_types"),
|
||||
pixel_values=batch.get("pixel_values"),
|
||||
image_grid_thw=batch.get("image_grid_thw"),
|
||||
proprioception=batch.get("proprioception"),
|
||||
dataset_names=batch.get("dataset_names"),
|
||||
dof_mask=batch.get("dof_mask"),
|
||||
agent_pos_mask=batch.get("agent_pos_mask"),
|
||||
prefix_length=batch.get("prefix_length"),
|
||||
positional_masks=batch.get("positional_masks"),
|
||||
re_generate=False,
|
||||
)
|
||||
|
||||
input_texts = model_output["input_text"]
|
||||
gt_texts = model_output["gt_output_text"]
|
||||
pred_texts = model_output["predict_output_text"]
|
||||
|
||||
if rank == 0:
|
||||
all_input_texts.extend(input_texts)
|
||||
all_gt_texts.extend(gt_texts)
|
||||
all_pred_texts.extend(pred_texts)
|
||||
for gt_text, pred_text in zip(gt_texts, pred_texts):
|
||||
gt_clean = gt_text[0] if isinstance(gt_text, list) else gt_text
|
||||
pred_clean = (
|
||||
pred_text[0] if isinstance(pred_text, list) else pred_text
|
||||
)
|
||||
point_l1_dist = calculate_point_l1_distance(gt_clean, pred_clean)
|
||||
if point_l1_dist is not None:
|
||||
all_point_l1_distances.append(point_l1_dist)
|
||||
|
||||
pepoch.update(1)
|
||||
|
||||
pepoch.close()
|
||||
|
||||
if rank == 0 and all_pred_texts:
|
||||
if all_point_l1_distances:
|
||||
avg_point_l1 = sum(all_point_l1_distances) / len(all_point_l1_distances)
|
||||
min_point_l1 = min(all_point_l1_distances)
|
||||
max_point_l1 = max(all_point_l1_distances)
|
||||
point_stats = {
|
||||
"text_prediction_point_l1_avg": avg_point_l1,
|
||||
"text_prediction_point_l1_min": min_point_l1,
|
||||
"text_prediction_point_l1_max": max_point_l1,
|
||||
"text_prediction_point_samples_count": len(all_point_l1_distances),
|
||||
}
|
||||
if logger is not None:
|
||||
logger.log(point_stats, step=current_step)
|
||||
if self.logger is not None:
|
||||
self.logger.info(
|
||||
f"Point L1 Statistics - Avg: {avg_point_l1:.4f}, "
|
||||
f"Min: {min_point_l1:.4f}, Max: {max_point_l1:.4f}, "
|
||||
f"Samples: {len(all_point_l1_distances)}"
|
||||
)
|
||||
save_text_results_to_file(
|
||||
current_step,
|
||||
all_input_texts,
|
||||
all_gt_texts,
|
||||
all_pred_texts,
|
||||
save_dir if save_dir is not None else "./",
|
||||
)
|
||||
|
||||
|
||||
def _move_batch(batch, device):
|
||||
"""Recursively move dict/list/tensor to device (VLA-local copy to keep
|
||||
VLAdapter self-contained; base_adapter has a similar helper used by
|
||||
validation)."""
|
||||
if isinstance(batch, dict):
|
||||
return {k: _move_batch(v, device) for k, v in batch.items()}
|
||||
if isinstance(batch, list):
|
||||
return [_move_batch(v, device) for v in batch]
|
||||
if torch.is_tensor(batch):
|
||||
return batch.to(device)
|
||||
return batch
|
||||
Reference in New Issue
Block a user