模型加载与初始化
importosimporthydraimporttorchfromomegaconfimportDictConfig,OmegaConffromhydra.utilsimportinstantiatefromtorch.utils.dataimportDataLoaderfromlightwam.runtimeimport_resolve_train_device,_normalize_mixed_precision,_mixed_precision_to_model_dtype,build_datasetsfromlightwam.utils.config_resolversimportregister_default_resolvers register_default_resolvers()defprint_tensor_info(name,tensor):ifisinstance(tensor,torch.Tensor):print(f" -{name}: shape={tuple(tensor.shape)}, dtype={tensor.dtype}, device={tensor.device}")else:print(f" -{name}: type={type(tensor).__name__}")@hydra.main(config_path="../configs",config_name="train",version_base="1.3")defmain(cfg:DictConfig):# --- 1. Initialize Device and Precision ---print("\n"+"="*50)print("1. Initialization Setup")print("="*50)model_device=_resolve_train_device()mixed_precision=_normalize_mixed_precision(cfg.mixed_precision)model_dtype=_mixed_precision_to_model_dtype(mixed_precision)print(f"Target Device:{model_device}")print(f"Mixed Precision:{mixed_precision}")print(f"Model Dtype:{model_dtype}")# --- 2. Model Instantiation ---print("\n"+"="*50)print("2. Model Initialization")print("="*50)print("Instantiating model from cfg.model...")# This calls lightwam.runtime.Wan22Runtime.from_config() under the hoodmodel=instantiate(cfg.model,model_dtype=model_dtype,device=model_device)print(f"Model Class:{type(model).__name__}")# If state_fusion is used, we can verify ituses_state_fusion=getattr(model,"uses_state_fusion_action_expert",lambda:False)()print(f"Uses State-Fusion Action Expert:{uses_state_fusion}")total_params=sum(p.numel()forpinmodel.parameters())print(f"Total Parameters:{total_params/1e6:.2f}M")# --- 3. Dataset Loading ---print("\n"+"="*50)print("3. Dataset Loading")print("="*50)print("Instantiating datasets from cfg.data...")train_ds,val_ds=build_datasets(cfg.data)print(f"Train Dataset:{type(train_ds).__name__}, Length:{len(train_ds)}")loader=DataLoader(train_ds,batch_size=int(cfg.batch_size),shuffle=False,num_workers=0,# single-threaded for quick testingpin_memory=False,)# Fetch exactly one batchprint("Fetching one batch from DataLoader...")batch=next(iter(loader))print(f"Batch Keys:{sorted(list(batch.keys()))}")# --- 4. Forward Pass (Loss Computation) ---print("\n"+"="*50)print("4. Model Forward Pass")print("="*50)# Put model in train modemodel.train()print("Moving batch to device and executing `model.training_loss(batch)`...")# We use autocast just like the trainer doeswithtorch.autocast(device_type=model_device.split(':')[0],dtype=model_dtype):# The model's `training_loss` internally handles moving relevant parts of `batch` to the correct deviceloss,loss_dict=model.training_loss(batch)print(f"\nForward pass successful!")print(f"Returned Total Loss:{loss.item():.4f}")print("Detailed Loss Dict:")fork,vinloss_dict.items():print(f" -{k}:{v:.4f}")print("\nDone. The script executed the exact flow used during training setup and first iteration.")if__name__=="__main__":main()模型运行结果
python scripts/inspect_train_flow.py\task=libero_uncond_2cam224_1e-4\data.train.dataset_dirs="['./data/libero_mujoco3.3.2/libero_goal_no_noops_lerobot']"\data.train.text_embedding_cache_dir='./data/text_embeds_cache/libero'\data.train.use_latent_cache=true\data.train.latent_cache_dir='./data/latent_cache_Wan2.1-T2V-1.3B/libero_goal_2cam224'\batch_size=2==================================================1. Initialization Setup==================================================Target Device: cuda:0 Mixed Precision: bf16 Model Dtype: torch.bfloat16==================================================2. Model Initialization==================================================Instantiating model from cfg.model...[2026-07-28 01:56:55,165][lightwam.models.wan22.helpers.loader][INFO]- Applying`wan2_1_t2v`video backbone preset overrides to`video_dit_config`:ffn_dim:14336->8960, hidden_dim:3072->1536, in_dim:48->16, num_heads:24->12, out_dim:48->16[2026-07-28 01:56:55,165][lightwam.models.wan22.helpers.loader][INFO]- Synchronized ActionDiT config with video backbone: num_heads:24->12[2026-07-28 01:56:55,165][lightwam.models.wan22.helpers.loader][INFO]- Loading Wan2.1-T2V-1.3B components...[2026-07-28 01:57:04,465][lightwam.models.wan22.wan_video_dit][INFO]- Enabled backbone LoRA onlayers=[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29]targets=['self_attn.q','self_attn.k','self_attn.v','self_attn.o','cross_attn.q','cross_attn.k','cross_attn.v','cross_attn.o','ffn.0','ffn.2']rank=64alpha=128.000dropout=0.000[2026-07-28 01:57:05,051][lightwam.models.wan22.helpers.loader][INFO]- Loaded WanVideoDiT from ./checkpoints/Wan-AI/Wan2.1-T2V-1.3B/diffusion_pytorch_model.safetensors using`raw`state dict variant(compatible=825,missing=618,unexpected=0).[2026-07-28 01:57:05,975][lightwam.models.wan22.helpers.loader][INFO]- Skipping pretrained text encoder/tokenizer load(`load_text_encoder=False`);training must provide cached`context/context_mask`.[2026-07-28 01:57:06,728][lightwam.models.wan22.helpers.loader][INFO]- Loaded WanVideoVAE from ./checkpoints/Wan-AI/Wan2.1-T2V-1.3B/Wan2.1_VAE.pth using`wan_video_vae_state_dict_converter`state dict variant(compatible=194,missing=0,unexpected=0).[2026-07-28 01:57:06,775][lightwam.models.wan22.helpers.loader][INFO]- Finished loading Wan2.1-T2V-1.3B componentsin11.61seconds.[2026-07-28 01:57:06,776][lightwam.models.wan22.mot][INFO]- Initialized MoT with experts:['video'],num_layers=30[2026-07-28 01:57:06,780][lightwam.models.wan22.mot][INFO]- Expert'video':num_params=1.51B Model Class: LightWAM Uses State-Fusion Action Expert: True Total Parameters:1986.82M==================================================3. Dataset Loading==================================================Instantiating datasets from cfg.data...[2026-07-28 01:57:09,334][datasets][INFO]- PyTorch version2.7.1+cu128 available. Resolving data files:100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████|433/433[00:00<00:00,25658.49it/s]Downloading data:100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████|433/433[00:00<00:00,50073.99files/s]Generating train split:52895examples[00:00,60523.66examples/s][2026-07-28 01:57:12,699][lightwam.datasets.lerobot.robot_video_dataset][INFO]- Calculating dataset statsfornormalization... Iterating dataset to get normalization:100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████|433/433[00:05<00:00,79.57it/s][2026-07-28 01:57:18,196][lightwam.datasets.lerobot.robot_video_dataset][INFO]- Loaded indexed latent cache index:format=sharded_v1shards=52samples=52895[2026-07-28 01:57:18,196][lightwam.datasets.lerobot.robot_video_dataset][INFO]- Using latent cacheforRobotVideoDataset: /workspace/Light-WAM/data/latent_cache_Wan2.1-T2V-1.3B/libero_goal_2cam224 Train Dataset: RobotVideoDataset, Length:52895Fetching one batch from DataLoader... Batch Keys:['action','action_is_pad','context','context_mask','idx','image_is_pad','prompt','proprio','proprio_is_pad','video_latents']==================================================4. Model Forward Pass==================================================Moving batch to device and executing`model.training_loss(batch)`... Forward pass successful!Returned Total Loss:1.4369Detailed Loss Dict: - loss_video:1.1475- loss_action:0.2894- loss_video_raw:1.1475- loss_action_raw:0.2894Done. The script executed the exact flow used during training setup and first iteration.