RuntimeError: The size of tensor a (8192) must match the size of tensor b (8189) at non-singleton dimension 2

**Describe the bug**

RuntimeError: The size of tensor a (8192) must match the size of tensor b (8189) at non-singleton dimension 2

**Steps/Code to reproduce bug**

Leverage yaml, python script, and test / validation files shared internally
Using Dtensor backend
Trying to run 70B on 8 H100 nodes

**Expected behavior**

Training should run based on the number of steps using 70B model

**Environment overview (please complete the following information)**

- Environment location: DGX Cloud Create Runai; GCP
- Method of install: git clone NeMo-RL w/ Dockerfile to build container from source
- If method of install is [Docker], provide `docker pull` & `docker run` commands used: N/A using kubernetes

**Environment details**

Using Nvidia provided Dockerfile

**Additional context**

deepseek-ai/DeepSeek-R1-Distill-Llama-70B
Zonal storage