模型加载与初始化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_nametrain,version_base1.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(fTarget Device:{model_device})print(fMixed Precision:{mixed_precision})print(fModel 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 hoodmodelinstantiate(cfg.model,model_dtypemodel_dtype,devicemodel_device)print(fModel Class:{type(model).__name__})# If state_fusion is used, we can verify ituses_state_fusiongetattr(model,uses_state_fusion_action_expert,lambda:False)()print(fUses State-Fusion Action Expert:{uses_state_fusion})total_paramssum(p.numel()forpinmodel.parameters())print(fTotal 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_dsbuild_datasets(cfg.data)print(fTrain Dataset:{type(train_ds).__name__}, Length:{len(train_ds)})loaderDataLoader(train_ds,batch_sizeint(cfg.batch_size),shuffleFalse,num_workers0,# single-threaded for quick testingpin_memoryFalse,)# Fetch exactly one batchprint(Fetching one batch from DataLoader...)batchnext(iter(loader))print(fBatch 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_typemodel_device.split(:)[0],dtypemodel_dtype):# The models training_loss internally handles moving relevant parts of batch to the correct deviceloss,loss_dictmodel.training_loss(batch)print(f\nForward pass successful!)print(fReturned 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\tasklibero_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_cachetrue\data.train.latent_cache_dir./data/latent_cache_Wan2.1-T2V-1.3B/libero_goal_2cam224\batch_size21. Initialization SetupTarget Device: cuda:0 Mixed Precision: bf16 Model Dtype: torch.bfloat162. Model InitializationInstantiating model from cfg.model...[2026-07-28 01:56:55,165][lightwam.models.wan22.helpers.loader][INFO]- Applyingwan2_1_t2vvideo backbone preset overrides tovideo_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]rank64alpha128.000dropout0.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 usingrawstate dict variant(compatible825,missing618,unexpected0).[2026-07-28 01:57:05,975][lightwam.models.wan22.helpers.loader][INFO]- Skipping pretrained text encoder/tokenizer load(load_text_encoderFalse);training must provide cachedcontext/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 usingwan_video_vae_state_dict_converterstate dict variant(compatible194,missing0,unexpected0).[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_layers30[2026-07-28 01:57:06,780][lightwam.models.wan22.mot][INFO]- Expertvideo:num_params1.51B Model Class: LightWAM Uses State-Fusion Action Expert: True Total Parameters:1986.82M3. Dataset LoadingInstantiating datasets from cfg.data...[2026-07-28 01:57:09,334][datasets][INFO]- PyTorch version2.7.1cu128 available. Resolving data files:100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████|433/433[00:0000:00,25658.49it/s]Downloading data:100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████|433/433[00:0000: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:0500:00,79.57it/s][2026-07-28 01:57:18,196][lightwam.datasets.lerobot.robot_video_dataset][INFO]- Loaded indexed latent cache index:formatsharded_v1shards52samples52895[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 PassMoving batch to device and executingmodel.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.