|
| 1 | +# Copyright (c) Meta Platforms, Inc. and affiliates. |
| 2 | +# All rights reserved. |
| 3 | +# |
| 4 | +# This source code is licensed under the BSD-style license found in the |
| 5 | +# LICENSE file in the root directory of this source tree. |
| 6 | + |
| 7 | +"""Characterization tests for Qwen3.5 position routing. |
| 8 | +
|
| 9 | +Locks the observable contract of the ``preprocess_inputs`` -> ``forward`` |
| 10 | +pipeline: for a folded (batch-flattened) dataloader batch, the position |
| 11 | +tensor that reaches every transformer layer is the ``(num_tokens, 3)`` |
| 12 | +``mrope_positions`` when present, else the plain 1D ``(num_tokens,)`` |
| 13 | +``positions`` -- while attention masks are always built from the 1D |
| 14 | +``positions``. |
| 15 | +
|
| 16 | +These go through the public ``preprocess_inputs`` seam and call |
| 17 | +``model(inputs, **batch)``, so they stay valid whether the mrope/positions |
| 18 | +resolution lives in ``forward`` or in ``preprocess_inputs``. |
| 19 | +""" |
| 20 | + |
| 21 | +import unittest |
| 22 | + |
| 23 | +import torch |
| 24 | +from torch import nn |
| 25 | + |
| 26 | + |
| 27 | +def _build_config_modules(): |
| 28 | + # torchtitan.models.qwen3_5 imports the FLA (flash-linear-attention) |
| 29 | + # kernels at module scope. FLA is a triton/CUDA-only optional dependency, |
| 30 | + # so skip instead of erroring on environments without it. |
| 31 | + try: |
| 32 | + from torchtitan.config import ParallelismConfig |
| 33 | + from torchtitan.distributed.parallel_dims import ParallelDims |
| 34 | + from torchtitan.models.qwen3_5 import model_registry |
| 35 | + except ModuleNotFoundError as exc: |
| 36 | + raise unittest.SkipTest( |
| 37 | + f"Qwen3.5 optional dependency unavailable: {exc.name}" |
| 38 | + ) from exc |
| 39 | + return model_registry, ParallelDims, ParallelismConfig |
| 40 | + |
| 41 | + |
| 42 | +class _RecordingLayer(nn.Module): |
| 43 | + """Layer stub that records the positions it is handed and passes x through. |
| 44 | +
|
| 45 | + The mrope/positions resolution is independent of the layer internals, so |
| 46 | + stubbing the layers keeps these tests on CPU without the FLA kernels while |
| 47 | + still exercising the real ``preprocess_inputs`` and ``forward`` glue. |
| 48 | + """ |
| 49 | + |
| 50 | + def __init__(self, sink: dict): |
| 51 | + super().__init__() |
| 52 | + self._sink = sink |
| 53 | + |
| 54 | + def forward(self, x, attention_masks=None, positions=None): |
| 55 | + self._sink["positions"] = positions |
| 56 | + return x |
| 57 | + |
| 58 | + |
| 59 | +class TestQwen35MRoPEPositions(unittest.TestCase): |
| 60 | + def _build_stub_model(self): |
| 61 | + model_registry, ParallelDims, ParallelismConfig = _build_config_modules() |
| 62 | + # varlen backend keeps mask construction to pure tensor ops (no flex |
| 63 | + # compile) so the pipeline runs on CPU. |
| 64 | + model = model_registry("debugmodel", attn_backend="varlen").model.build() |
| 65 | + sink: dict = {} |
| 66 | + for key in list(model.layers.keys()): |
| 67 | + model.layers[key] = _RecordingLayer(sink) |
| 68 | + parallel_dims = ParallelDims( |
| 69 | + dp_replicate=1, dp_shard=1, cp=1, tp=1, pp=1, ep=1, world_size=1 |
| 70 | + ) |
| 71 | + # partial_dtensor avoids the spmd_types annotation path, which is |
| 72 | + # orthogonal to position routing. |
| 73 | + parallelism = ParallelismConfig(spmd_backend="partial_dtensor") |
| 74 | + return model, sink, parallel_dims, parallelism |
| 75 | + |
| 76 | + def _run(self, model, parallel_dims, parallelism, input_dict): |
| 77 | + inputs, _labels, batch = model.preprocess_inputs( |
| 78 | + input_dict, |
| 79 | + parallel_dims=parallel_dims, |
| 80 | + parallelism=parallelism, |
| 81 | + ) |
| 82 | + model(inputs, **batch) |
| 83 | + return batch |
| 84 | + |
| 85 | + def test_text_batch_routes_1d_positions_to_layers(self): |
| 86 | + model, sink, parallel_dims, parallelism = self._build_stub_model() |
| 87 | + # Folded 1D token stream packing docs of length 3, 2, and 5. |
| 88 | + positions = torch.tensor([0, 1, 2, 0, 1, 0, 1, 2, 3, 4], dtype=torch.int32) |
| 89 | + input_dict = { |
| 90 | + "input": torch.randint(0, 100, (10,)), |
| 91 | + "positions": positions.clone(), |
| 92 | + "labels": torch.zeros(10), |
| 93 | + } |
| 94 | + |
| 95 | + batch = self._run(model, parallel_dims, parallelism, input_dict) |
| 96 | + |
| 97 | + # No mrope: layers see the plain 1D positions. |
| 98 | + self.assertTrue(torch.equal(sink["positions"], positions)) |
| 99 | + # Masks come from the 1D positions. |
| 100 | + self.assertEqual( |
| 101 | + batch["attention_masks"]["deltanet"].cu_seq_q_host, (0, 3, 5, 10) |
| 102 | + ) |
| 103 | + |
| 104 | + def test_multimodal_batch_routes_mrope_to_layers(self): |
| 105 | + model, sink, parallel_dims, parallelism = self._build_stub_model() |
| 106 | + positions = torch.tensor([0, 1, 2, 0, 1, 0, 1, 2, 3, 4], dtype=torch.int32) |
| 107 | + # Folded (num_tokens, 3) T/H/W positions whose H/W channels differ from |
| 108 | + # the 1D positions, so routing the wrong tensor to the layers is |
| 109 | + # detectable. |
| 110 | + mrope_positions = torch.stack( |
| 111 | + [positions, positions + 7, positions + 13], dim=-1 |
| 112 | + ) |
| 113 | + input_dict = { |
| 114 | + "input": torch.randint(0, 100, (10,)), |
| 115 | + "positions": positions.clone(), |
| 116 | + "mrope_positions": mrope_positions.clone(), |
| 117 | + "labels": torch.zeros(10), |
| 118 | + } |
| 119 | + |
| 120 | + batch = self._run(model, parallel_dims, parallelism, input_dict) |
| 121 | + |
| 122 | + # mrope present: layers see the (num_tokens, 3) mrope positions, not the |
| 123 | + # 1D positions. |
| 124 | + self.assertEqual(sink["positions"].ndim, 2) |
| 125 | + self.assertEqual(sink["positions"].shape[-1], 3) |
| 126 | + self.assertTrue(torch.equal(sink["positions"], mrope_positions)) |
| 127 | + # Masks are still built from the 1D positions, not the mrope positions. |
| 128 | + self.assertEqual( |
| 129 | + batch["attention_masks"]["deltanet"].cu_seq_q_host, (0, 3, 5, 10) |
| 130 | + ) |
| 131 | + |
| 132 | + |
| 133 | +if __name__ == "__main__": |
| 134 | + unittest.main() |
0 commit comments