SolitaryThinker
d81f1f166b
- MiniMaxH3Model: allow per-role attention backends (TORCH_SDPA/FLASH_ATTN/ VIDEO_SPARSE_ATTN_H3); predict_noise routes attn_metadata_vsa for attn_kind='vsa' (dense roles keep None). - MiniMaxH3DMDModel.prepare_batch builds real VSA-H3 metadata (packed text|cond|audio prefix segments via _h3_vsa_prefix_segments, sparsity from training.vsa.sparsity) when the role runs VIDEO_SPARSE_ATTN_H3. - VSA-H3 tile(): fresh buffer for grad-tracking forwards; shared-buffer reuse trips autograd's in-place version check at student backward. - build_tracker: degrade to DummyTracker when wandb is unimportable or has no credentials (modular stack already logs per-step loss dicts via the wandb tracker; no new callback needed). - dmd2_vsa0_overfit.yaml: student VSA-H3 @ sparsity 0.0, teacher/critic FLASH_ATTN, 3-step DMD schedule, wandb project h3-dmd2-vsa. - preprocess_minimax_h3_overfit: --video/--prompt/--output-dir/--model-path CLI; overfit_vsa0.sh: preprocess + 4-GPU torchrun with FASTVIDEO_FA4=1. - CPU tests: VSA metadata build/routing, per-role backend override, wandb guard + loss-dict logging, tile-under-grad regression.
79 lines
2.3 KiB
Python
79 lines
2.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from typing import Any, TYPE_CHECKING
|
|
|
|
from fastvideo.distributed import get_world_group
|
|
from fastvideo.logger import init_logger
|
|
from fastvideo.training.trackers import (
|
|
initialize_trackers,
|
|
Trackers,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from fastvideo.train.utils.training_config import (
|
|
CheckpointConfig,
|
|
TrackerConfig,
|
|
)
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
def _wandb_usable() -> bool:
|
|
"""True when wandb can actually start a run (importable + credentials)."""
|
|
try:
|
|
import wandb
|
|
except Exception:
|
|
return False
|
|
if os.environ.get("WANDB_API_KEY"):
|
|
return True
|
|
if os.environ.get("WANDB_MODE") in ("offline", "disabled"):
|
|
return True
|
|
try:
|
|
# Covers ~/.netrc logins from a prior `wandb login`.
|
|
return wandb.api.api_key is not None
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
def build_tracker(
|
|
tracker_config: TrackerConfig,
|
|
checkpoint_config: CheckpointConfig,
|
|
*,
|
|
config: dict[str, Any] | None,
|
|
) -> Any:
|
|
"""Build a tracker instance for a distillation run."""
|
|
|
|
world_group = get_world_group()
|
|
|
|
trackers = list(tracker_config.trackers)
|
|
if not trackers and str(tracker_config.project_name):
|
|
trackers.append(Trackers.WANDB.value)
|
|
if world_group.rank != 0:
|
|
trackers = []
|
|
if Trackers.WANDB.value in trackers and not _wandb_usable():
|
|
logger.warning("wandb tracking requested but wandb is not importable "
|
|
"or no credentials are configured (WANDB_API_KEY); "
|
|
"continuing without the wandb tracker.")
|
|
trackers = [t for t in trackers if t != Trackers.WANDB.value]
|
|
|
|
tracker_log_dir = (checkpoint_config.output_dir or os.getcwd())
|
|
if trackers:
|
|
tracker_log_dir = os.path.join(tracker_log_dir, "tracker")
|
|
|
|
tracker_config_dict = config if trackers else None
|
|
tracker_entity = tracker_config.entity or None
|
|
tracker_run_name = tracker_config.run_name or None
|
|
project = (tracker_config.project_name or "fastvideo")
|
|
|
|
return initialize_trackers(
|
|
trackers,
|
|
experiment_name=project,
|
|
config=tracker_config_dict,
|
|
log_dir=tracker_log_dir,
|
|
entity=tracker_entity,
|
|
run_name=tracker_run_name,
|
|
)
|