[Bug] TypeError when training hunyuan model due to NoneType mask_strategy

[dingangui]

Environment

A800*4

Describe the bug

Issue Description:

When attempting to train the hunyuan model using the script scripts/finetune/finetune_hunyuan.sh at commit verison 1e08893, I encountered a TypeError:

TypeError: 'NoneType' object is not subscriptable

This error occurred at the following code location fastvideo/models/hunyuan/modules/models.py due to mask_strategy[index]

for index, block in enumerate(self.double_blocks):
    double_block_args = [img, txt, vec, freqs_cis, text_mask, mask_strategy[index]]
    img, txt = block(*double_block_args)

when the code runs at fastvideo/train.py:

        with torch.autocast("cuda", dtype=torch.bfloat16):
            input_kwargs = {
                "hidden_states": noisy_model_input,
                "encoder_hidden_states": encoder_hidden_states,
                "timestep": timesteps,
                "encoder_attention_mask": encoder_attention_mask,  # B, L
                "return_dict": False,
            }
            if 'hunyuan' in model_type:
                input_kwargs["guidance"] = torch.tensor([1000.0], device=noisy_model_input.device, dtype=torch.bfloat16)
            model_pred = transformer(**input_kwargs)[0]

Problem Analysis:

The error arises because mask_strategy is None, mask_strategy is not specified in input_kwargs, making it impossible to perform the indexing operation mask_strategy[index].

Upon reviewing the commit history, I noticed that this issue was addressed in a subsequent commit 0be4fc6 with the following fix:

        if mask_strategy is None:
            mask_strategy = [[None] * len(self.heads_num)
                             for _ in range(len(self.double_blocks) + len(self.single_blocks))]

However, this fix introduces a new issue: self.heads_num is an integer type and does not support the len() function. This creates a potential bug waiting to happen.

Questions:

How should I properly initialize mask_strategy to prevent such errors? and by the way, what is mask_stratage for?

Reproduction

bash scripts/finetune/finetune_hunyuan.sh

[BrianChen1129]

Thanks for pointing out this. It should be fixed now! Mask_strategy is used for STA inference.

[yinian-lw]

same problem