attention mask fixes

[ahmadki]

What does this PR do ?

Add a one line overview of what this PR aims to accomplish.

Issues

List issues that this PR closes (syntax):

Usage

• You can potentially add a usage example below

# Add a code snippet demonstrating how to use this 

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

• ...

[terrykong]

@ahmadki can you rebase? I think the dtensor attention mask diff is already resolved top-of-tree. The fsdp1 makes sense. nice find

[ahmadki]

I rebased with main.

I don't think the issue with dtensor is resolved in main. attention_mask is created but never used in the train function. Instead, we use an all ones tensor instead here: https://github.com/NVIDIA/nemo-rl/blob/ebb46c3b936e6c31494dae3c7f0953bfeff006fb/nemo_rl/models/policy/dtensor_policy_worker.py#L316

I'm not sure if attention_mask should be removed completely, or it should be taken into account when calculating the logprobs here: https://github.com/NVIDIA/nemo-rl/blob/ebb46c3b936e6c31494dae3c7f0953bfeff006fb/nemo_rl/models/policy/dtensor_policy_worker.py#L332
Either ways, someone needs to take a look at it.