Diff Coverage

    Composition pattern — delegates training loop to BaseTrainer.
    """

    def __init__(self, args):
        self.base = BaseTrainer(args)

        # 13-step init — call base methods, override data steps
        self.base._setup()
        self.base._build_model()
        # 注意：不要手动 .to(device)，_build_parallelized_model 会处理 meta → npu
        self.base._freeze_model()
        self._build_model_assets()
        self._build_data_transform()
        self._build_dataset()
        self._build_collate_fn()
        self._build_dataloader()
        self.base._build_parallelized_model()
        self._build_optimizer()
        self.base._build_lr_scheduler()
        self.base._build_training_context()
        self.base._init_callbacks()
        self.base.on_init_end()

    # ------------------------------------------------------------------
    # Overridden _build_* methods
    # ------------------------------------------------------------------

    # ------------------------------------------------------------------

    def _build_model_assets(self):
        """DiT does not need tokenizer or processor for dummy data."""
        self.base.tokenizer = None
        self.base.processor = None

    def _build_data_transform(self):
        """No offline data transform for dummy tensors."""
        self.base.data_transform = None

    def _build_dataset(self):
        """Build deterministic dummy DiT dataset with real diffusion targets.

        base_seed = int(getattr(self.base.args, "seed", 42))
        max_steps = getattr(self.base.args.train, "max_steps", 100)

        if data_type == "coco_parquet":
            self._build_parquet_dataset(max_steps)
        elif data_type == "coco_dit":
            self._build_coco_dataset(cond_dim, seq_len, base_seed, max_steps)
        elif data_type == "dummy_dit":
            self._build_dummy_dataset(model_cfg, cond_dim, seq_len, base_seed, max_steps)
        else:
            raise NotImplementedError(
                f"DiTTrainer supports 'dummy_dit'/'coco_dit'/'coco_parquet', got '{data_type}'"
            )

        )

    def _build_coco_dataset(self, cond_dim, seq_len, base_seed, max_steps):
        """Real COCO images (packed VAE latents) + dummy text embeddings."""
        cache_path = getattr(self.base.args.data, "train_path", None)
        if not cache_path:
            raise ValueError("data.train_path must be set for data.type=coco_dit")
        pth_files = sorted(glob.glob(os.path.join(cache_path, "*.pth")))
        if not pth_files:
            raise FileNotFoundError(f"No .pth files found in {cache_path}")

        class CocoDiTDataset(Dataset):
            """Map-style dataset over packed-VAE .pth files with on-the-fly
            flow-matching noise targets."""
            def __init__(self, files, cond_dim, seq_len, seed):
                self.files = files
                self.cond_dim = cond_dim
                self.seq_len = seq_len
                self.seed = seed

            def __len__(self):
                return len(self.files)

            def __getitem__(self, idx):
                data = torch.load(self.files[idx])
                clean = data["latent_clean"]
                g = torch.Generator().manual_seed(self.seed + idx)
                eps = torch.randn(*clean.shape, generator=g)
                ts = torch.randint(1, 1000, (1,), generator=g).squeeze(0)
                t_norm = ts.float() / 1000.0
                x_t = (1.0 - t_norm) * clean + t_norm * eps
                velocity = eps - clean
                if "text_embed" in data and data["text_embed"] is not None:
                    condition = data["text_embed"]
                else:
                    condition = torch.randn(self.seq_len, self.cond_dim, generator=g)
                return {
                    "latent": x_t, "timestep": ts,
                    "condition": condition, "target_noise": velocity,
                    "labels": torch.tensor(1, dtype=torch.long),
                }

                    "condition": condition, "target_noise": velocity,
                    "labels": torch.tensor(1, dtype=torch.long),
                }

        self.base.train_dataset = CocoDiTDataset(pth_files, cond_dim, seq_len, base_seed)
        self.base.state.max_steps = max_steps
        logger.info_rank0(
            f"COCO dataset: {len(pth_files)} samples from {cache_path}, "
            f"cond=({seq_len},{cond_dim})"
        )

        Both HP and VeOmni load the same parquet through HuggingFace Datasets,
        ensuring byte-identical training inputs for cross-framework alignment.
        """
        try:
            from datasets import load_dataset  # pylint: disable=import-outside-toplevel
            import pickle as pk  # pylint: disable=import-outside-toplevel
        except ImportError as exc:
            raise ImportError("datasets package required: pip install datasets") from exc

        parquet_path = getattr(self.base.args.data, "train_path", None)
        if not parquet_path:
            raise ValueError("data.train_path must point to a .parquet file for data.type=coco_parquet")

        hf_ds = load_dataset("parquet", data_files=parquet_path, split="train")

        class ParquetDataset(Dataset):
            """Map-style dataset reading pre-generated rows from a parquet file."""
            def __init__(self, hf_ds):
                self.hf_ds = hf_ds
            def __len__(self):
                return len(self.hf_ds)
            def __getitem__(self, idx):
                row = self.hf_ds[idx]
                hidden = pk.loads(row["hidden_states"]).squeeze(0)
                target = pk.loads(row["training_target"]).squeeze(0)
                c, h, w = 64, 16, 16
                return {
                    "latent": hidden.reshape(h, w, c).permute(2, 0, 1).float(),
                    "timestep": pk.loads(row["timestep"]).squeeze(0),
                    "condition": pk.loads(row["encoder_hidden_states"]).squeeze(0).float(),
                    "target_noise": target.reshape(h, w, c).permute(2, 0, 1).float(),

                    "target_noise": target.reshape(h, w, c).permute(2, 0, 1).float(),
                    "labels": torch.tensor(1, dtype=torch.long),
                }

        self.base.train_dataset = ParquetDataset(hf_ds)
        self.base.state.max_steps = max_steps
        logger.info_rank0(
            f"Parquet dataset: {len(hf_ds)} samples from {parquet_path}"
        )

    def _build_collate_fn(self):

    def _build_collate_fn(self):
        """Stack fixed-size tensors (no padding needed for dummy data)."""

        def _dit_collate(batch):
            return {
                "latent": torch.stack([x["latent"] for x in batch]),
                "timestep": torch.stack([x["timestep"] for x in batch]),
                "condition": torch.stack([x["condition"] for x in batch]),
                "target_noise": torch.stack([x["target_noise"] for x in batch]),

                "target_noise": torch.stack([x["target_noise"] for x in batch]),
                "labels": torch.stack([x["labels"] for x in batch]),
            }

        self.base.collate_fn = _dit_collate

    def _build_dataloader(self):
        """Build DataLoader with no distributed sharding.

        no_decay_params = []
        seen_ids = set()
        for n, p in self.base.model.named_parameters():
            if not p.requires_grad:
                continue
            if id(p) in seen_ids:
                continue
            seen_ids.add(id(p))
            if _is_no_decay(n):
                no_decay_params.append(p)
            else:

    # ------------------------------------------------------------------

    def train(self):
        """Delegate to BaseTrainer.train()."""
        return self.base.train()