fixed validation error when using flash attention

[francesco-bertolotti]

This PR addresses the issue #2140

Briefly, the bug is related to the validation step.
The validate method did not pass the attention mask to the model.
When using flash attention, this leads to an error.

[wwwjn]

Thanks for the fix! Please fix the PP enabled branch as well, I think we would do something similar to train.py: https://github.com/pytorch/torchtitan/blob/refs/heads/main/torchtitan/train.py#L507-L516

[wwwjn]

For readability, can you make this into a function similar as: https://github.com/pytorch/torchtitan/blob/refs/heads/main/torchtitan/train.py#L416C9-L416C33

[francesco-bertolotti]

The new commit propagated the fix to the pipeline parallelism case. Further, I have encapsulated the creation of the attention mask in a method.

[tianyu-l]

can we reuse the one in train.py?

[francesco-bertolotti]

In principle, it would be possible, but I would need to make some modifications since their signatures differ. In particular, the trainer accesses model_parts through self.

One possible solution is to turn post_dataloading_process into a utility function that both the trainer and validator can call. However, this would mean you can no longer modify the behavior through inheritance—unless the current post_dataloading_process simply becomes a wrapper around the new utility function.|

If you have any suggestions, I can add a commit.

[tianyu-l]

I see.
If you want to "modify the behavior", we probably should have a PostDataloadingProcessor class.
If not, I think it's fine to call a util directly in both trainer and validator.

WDYT?
cc @fegin

[fegin]

I'm thinking this as well, I'm seeing the duplicated code in trainer and validator, model_parts is not a big problem, you can add it. See #2144. But the duplicated logic is pretty concerning.

My initial proposal is to have an util function in distributed.utils so that both methods can call it. The only worry is that I don't know whether distributed.utils is the best place to put it. It is not now because the util function purely unwrap the input but the util function will encapsulate the CP logic after #2144. So distributed.utils seems to be a good place.

@francesco-bertolotti you can do it as your PR may be landed first. I can rebase on top of your change.

[tianyu-l]

agreed we should share code, but not sure about distributed.utils. Let's discuss offline.

[tianyu-l]

unblocking

[tianyu-l]

@francesco-bertolotti CPU unit test failed, could you take a look?

[francesco-bertolotti]

Hi @tianyu-l,

I looked into the error, and it turns out to be a circular import issue.

In validate.py, I need access to model_args. Since Validator doesn’t have direct access to it, I currently construct it via get_train_spec. However, importing get_train_spec triggers an import of BaseValidator, which lives in the same module as Validator, leading to the circular dependency.

I see two possible ways to resolve this:

1. Pass model_args directly to Validator, eliminating the need to import get_train_spec.
2. Move BaseValidator into its own module.

I’m leaning slightly toward changing the Validator signature, as it feels cleaner overall, though it would be a breaking change.

Curious to hear your thoughts.

[tianyu-l]

It seems this is being fixed in #2144, with the first approach you mentioned.

Are you OK if we consolidate effort there?