fix: do not pass positions to RoPE when attn_backend is sdpa

Dockerfile

@@ -0,0 +1,72 @@
+FROM rocm/pytorch-nightly:2026-02-17-rocm7.2
+
+RUN apt-get update && apt-get install -y \
+    git-lfs \
+    pkg-config \
+    clang \
+    libclang-dev \
+    libunwind-dev \
+    libnl-3-dev \
+    libnl-route-3-dev \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN update-pciids
+
+RUN pip install --no-cache-dir huggingface_hub "datasets>=3.6.0" \
+    transformers tabulate wandb fsspec tyro "tokenizers>=0.15.0" safetensors \
+    tensorboard pre-commit yapf pybind11 meson-python torchdata pytablewriter \
+    "antlr4-python3-runtime==4.11.0" sympy math_verify more_itertools peft \
+    accelerate pillow "numpy<2" opencv-python-headless scipy \
+    numba huggingface-hub[cli,hf_transfer] "packaging>=24.2" \
+    "setuptools>=77.0.3,<80.0.0" "setuptools-scm>=8" \
+    protobuf-protoc-bin fmt && \
+    pip install --no-cache-dir /opt/rocm/share/amd_smi
+
+RUN cd /var/lib/jenkins && \
+    git clone --depth 1 https://github.com/EleutherAI/lm-evaluation-harness && \
+    cd lm-evaluation-harness && \
+    pip install -e .
+
+ENV PYTORCH_ROCM_ARCH="gfx90a;gfx942;gfx950"
+ENV RUSTUP_HOME=/opt/rustup
+ENV CARGO_HOME=/opt/cargo
+
+RUN rm -rf /root/.rustup && \
+    curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
+
+RUN cd /var/lib/jenkins && \
+    git clone --recursive https://github.com/linux-rdma/rdma-core.git && \
+    cd rdma-core && \
+    git checkout v62.0 && \
+    mkdir build && \
+    cd build && \
+    cmake -DNO_MAN_PAGES=1 .. && \
+    make -j$(nproc) && \
+    make install && \
+    ldconfig
+
+RUN cd /var/lib/jenkins && \
+    . $CARGO_HOME/env && \
+    git clone --recursive https://github.com/hann-wang/monarch.git && \
+    cd monarch && \
+    git checkout han/rocm && \
+    pip install -r build-requirements.txt && \
+    USE_TENSOR_ENGINE=1 \
+    pip install --no-build-isolation -e .
+
+RUN cd /var/lib/jenkins && \
+    git clone https://github.com/vllm-project/vllm.git && \
+    cd vllm && \
+    git checkout v0.17.0 && \
+    sed -i 's/^opencv-python-headless.*//' requirements/common.txt && \
+    sed -i 's/^ray.*//' requirements/common.txt && \
+    pip install -r requirements/rocm.txt "numpy<2" && \
+    python3 setup.py develop
+
+# COPY . /usr/local/src/torchtitan
+
+# RUN cd /usr/local/src/torchtitan && \
+#     pip install --no-build-isolation --no-deps -e .
+
+RUN sed -i 's/self.sharded_param = nn.Parameter(self.to_sharded_dtensor(sharded_param))/self.sharded_param = nn.Parameter(self.to_sharded_dtensor(sharded_param), requires_grad=param.requires_grad)/' /opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/distributed/fsdp/_fully_shard/_fsdp_param.py && \
+    sed -i 's/        self.sharded_param.requires_grad_(param.requires_grad)//' /opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/distributed/fsdp/_fully_shard/_fsdp_param.py

docker/build.sh

@@ -0,0 +1,8 @@
+#!/bin/bash
+
+IMAGE=ghcr.io/amd-agi/han-workspace:ubuntu22.04-pytorch2.12.0dev20260217-rocm7.2
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
+cd $SCRIPT_DIR/..
+
+docker build -t $IMAGE .
+docker push $IMAGE

docker/run.sh

@@ -0,0 +1,19 @@
+#!/bin/bash
+
+IMAGE=ghcr.io/amd-agi/han-workspace:ubuntu22.04-pytorch2.12.0dev20260217-rocm7.2-patch
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
+GID_RENDER=$(getent group render | cut -d: -f3)
+GID_VIDEO=$(getent group video | cut -d: -f3)
+
+docker run --rm -it -u $(id -u):$GID_RENDER \
+    --ulimit core=0  --privileged \
+    --cap-add=SYS_PTRACE --security-opt seccomp=unconfined \
+    --device=/dev/kfd \
+    --device=/dev/dri \
+    --group-add $GID_RENDER \
+    --group-add $GID_VIDEO \
+    --network host \
+    --ipc=host --shm-size 8G \
+    --workdir /workspace \
+    -v $SCRIPT_DIR/../..:/workspace \
+    $IMAGE

meson.build

@@ -0,0 +1,32 @@
+project(
+    'torchtitan',
+    'cpp',
+    version: run_command('cat', 'assets/version.txt', check: true).stdout().strip(),
+    default_options: [
+        'optimization=3',
+        'warning_level=1',
+        'cpp_std=c++17',
+    ],
+)
+
+find = find_program('find')
+res = run_command(find, 'torchtitan', '-type', 'f', '-name', '*.py', '-or', '-name', '*.safetensors')
+lines = res.stdout().strip().split('\n')
+sources = []
+foreach f : lines
+    if f != ''
+        sources += f
+    endif
+endforeach
+
+py = import('python').find_installation(pure: false)
+pybind11_dep = dependency('pybind11', required: true)
+
+py.install_sources(sources, preserve_path: true)
+
+py.extension_module('helpers_cpp',
+    ['torchtitan/hf_datasets/helpers.cpp'],
+    install: true,
+    subdir: 'torchtitan/hf_datasets',
+    dependencies : [py.dependency(), pybind11_dep],
+)

pyproject.toml

@@ -42,21 +42,14 @@ dev = [
     "pyrefly==0.45.1",
 ]
 
-[tool.setuptools.dynamic]
-version = {file = "assets/version.txt"}
-
 [tool.codespell]
 skip = ["*.json"]
 ignore-words-list = ["assertin", "datas" ,"indx", "inpt", "nd", "socio-economic"]
 
 # ---- Explicit project build information ---- #
 [build-system]
-requires = ["setuptools>=61.0"]
-build-backend = "setuptools.build_meta"
-
-[tool.setuptools.packages.find]
-where = [""]
-include = ["torchtitan*"]
+requires = ["meson-python", "pybind11"]
+build-backend = "mesonpy"
 
 [tool.pytest.ini_options]
 addopts = ["--showlocals"]  # show local variables in tracebacks

tests/unit_tests/test_sdpa_rope_positions.py

@@ -0,0 +1,68 @@
+import unittest
+from unittest.mock import patch
+
+import torch
+
+from torchtitan.models.common.attention import GQAttention
+
+
+class _StopAfterRoPE(Exception):
+    pass
+
+
+class TestGQAttentionRoPEPositions(unittest.TestCase):
+    def _run_until_rope(self, attn_backend: str, rope_backend: str):
+        attention = GQAttention(
+            GQAttention.Config(
+                n_heads=2,
+                n_kv_heads=1,
+                attn_backend=attn_backend,
+                rope_backend=rope_backend,
+            ),
+            dim=8,
+        )
+        positions = torch.tensor([[4, 5, 6]])
+        captured_positions = []
+
+        def stop_after_rope(*args, **kwargs):
+            received_positions = (
+                kwargs["positions"] if "positions" in kwargs else args[3]
+            )
+            captured_positions.append(received_positions)
+            raise _StopAfterRoPE
+
+        rope_function = (
+            "apply_rotary_emb_cos_sin"
+            if rope_backend == "cos_sin"
+            else "apply_rotary_emb_complex"
+        )
+        with patch(
+            f"torchtitan.models.common.attention.{rope_function}",
+            side_effect=stop_after_rope,
+        ):
+            with self.assertRaises(_StopAfterRoPE):
+                attention(
+                    torch.randn(1, 3, 8),
+                    torch.randn(8, 2),
+                    attention_masks=None,
+                    positions=positions,
+                )
+
+        return captured_positions[0], positions
+
+    def test_sdpa_drops_positions_for_both_rope_formats(self):
+        for rope_backend in ("complex", "cos_sin"):
+            with self.subTest(rope_backend=rope_backend):
+                received_positions, _ = self._run_until_rope("sdpa", rope_backend)
+                self.assertIsNone(received_positions)
+
+    def test_non_sdpa_keeps_positions_for_both_rope_formats(self):
+        for attn_backend in ("flex", "varlen"):
+            for rope_backend in ("complex", "cos_sin"):
+                with self.subTest(
+                    attn_backend=attn_backend, rope_backend=rope_backend
+                ):
+                    received_positions, positions = self._run_until_rope(
+                        attn_backend, rope_backend
+                    )
+                    self.assertIs(received_positions, positions)

torchtitan/components/checkpoint.py

@@ -22,7 +22,7 @@
 import torch.distributed as dist
 import torch.distributed.checkpoint as dcp
 import torch.nn as nn
-from torch.distributed.checkpoint import HuggingFaceStorageWriter
+from torch.distributed.checkpoint import DefaultLoadPlanner, HuggingFaceStorageWriter
 from torch.distributed.checkpoint._consolidate_hf_safetensors import (
     consolidate_safetensors_files_on_every_rank,
 )
@@ -347,6 +347,7 @@ class Config(Configurable.Config):
     mp_queue_send: queue.Queue
     pg: dist.ProcessGroup
     purge_thread: threading.Thread | None
+    allow_partial_load: bool = False
 
     def __init__(
         self,
@@ -620,10 +621,16 @@ def dcp_load(
             hf_storage_reader = self.sd_adapter.get_hf_storage_reader(
                 checkpoint_id, from_quantized
             )
+            planner = DefaultLoadPlanner(
+                flatten_state_dict=True,
+                flatten_sharded_tensors=True,
+                allow_partial_load=self.allow_partial_load,
+            )
 
             dcp.load(
                 hf_state_dict,
                 storage_reader=hf_storage_reader,
+                planner=planner,
             )
 
             state_dict = self.sd_adapter.from_hf(hf_state_dict)

torchtitan/components/validate.py

@@ -92,7 +92,6 @@ def __post_init__(self):
                 self.steps > 0 or self.steps == -1
             ), "validation steps must be positive or -1"
 
-    validation_dataloader: BaseDataLoader
 
     # TODO: improve the constructor signature
     def __init__(
@@ -120,15 +119,11 @@ def __init__(
         self.tokenizer = tokenizer
         self.parallel_dims = parallel_dims
         self.loss_fn = loss_fn
-        # pyrefly: ignore [unexpected-keyword]
-        dl_config = replace(config.dataloader, infinite=config.steps != -1)
-        self.validation_dataloader = dl_config.build(
-            dp_world_size=dp_world_size,
-            dp_rank=dp_rank,
-            tokenizer=tokenizer,
-            seq_len=seq_len,
-            local_batch_size=local_batch_size,
-        )
+        self.dl_config = replace(config.dataloader, infinite=config.steps != -1)
+        self.dp_world_size = dp_world_size
+        self.dp_rank = dp_rank
+        self.seq_len = seq_len
+        self.local_batch_size = local_batch_size
         self.validation_context = validation_context
         self.maybe_enable_amp = maybe_enable_amp
         self.metrics_processor = metrics_processor
@@ -225,7 +220,15 @@ def validate(
         device_type = utils.device_type
         num_steps = 0
 
-        for input_dict, labels in self.validation_dataloader:
+        validation_dataloader = self.dl_config.build(
+            dp_world_size=self.dp_world_size,
+            dp_rank=self.dp_rank,
+            tokenizer=self.tokenizer,
+            seq_len=self.seq_len,
+            local_batch_size=self.local_batch_size,
+        )
+
+        for input_dict, labels in validation_dataloader:
             # pyrefly: ignore [missing-attribute, unsupported-operation]
             if self.config.steps != -1 and num_steps >= self.config.steps:
                 break

torchtitan/distributed/tensor_parallel.py

@@ -134,7 +134,7 @@ def __init__(
         self.local_input_grad_placements = local_input_grad_placements
 
     @staticmethod
-    def _prepare_input_fn(  # pyrefly: ignore [bad-param-name-override]
+    def _prepare_input_fn(  # pyrefly: ignore [bad-override]
         input_layouts,
         desired_input_layouts,
         local_input_grad_placements,

torchtitan/experiments/forge/engine.py

@@ -159,6 +159,16 @@ def __init__(self, config: Config):
         ):
             model = model_config.build()
 
+        # Build the collection of model converters. No-op if converters empty
+        model_compile_enabled = (
+            config.compile.enable and "model" in config.compile.components
+        )
+        self.model_converters = config.model_converters.build(
+            parallel_dims=parallel_dims,
+            model_compile_enabled=model_compile_enabled,
+        )
+        self.model_converters.convert(model)
+
         # calculate model size and flops per token
         (
             self.model_param_count,
@@ -256,6 +266,8 @@ def __init__(self, config: Config):
 
             self.model_parts = [model]
 
+        self.model_converters.post_initialization(self.model_parts)
+
         # build optimizer after applying parallelisms to the model
         self.optimizers = config.optimizer.build(
             model_parts=self.model_parts,
@@ -268,6 +280,14 @@ def __init__(self, config: Config):
             optimizers=self.optimizers,
             training_steps=config.training.steps,
         )
+        # Post optimizer step model converters hook.
+        # e.g. calculate float8 dynamic amax/scale for all-parameter for FSDP2
+        # where it issues a single all-reduce for all parameters at once for better performance
+        self.optimizers.register_step_post_hook(
+            lambda *args, **kwargs: self.model_converters.post_optimizer_hook(
+                self.model_parts
+            )
+        )
 
         self.checkpointer = config.checkpoint.build(
             dataloader=None,