From 2cd4e843a0f01f6174b34ada8e6452f6f031d512 Mon Sep 17 00:00:00 2001 From: Kazuaki Ishizaki Date: Wed, 22 Jul 2026 02:28:28 +0000 Subject: [PATCH 1/4] use xavier for initialization Signed-off-by: Kazuaki Ishizaki --- .../auto_generate_module_config.py | 63 +++++++++++++++++-- 1 file changed, 59 insertions(+), 4 deletions(-) diff --git a/utils/module_discovery/auto_generate_module_config.py b/utils/module_discovery/auto_generate_module_config.py index a860ba6a..d7add9e5 100644 --- a/utils/module_discovery/auto_generate_module_config.py +++ b/utils/module_discovery/auto_generate_module_config.py @@ -83,6 +83,21 @@ def _is_special_tensor(name: str) -> bool: # config faithful to the runtime implementation. DEFAULT_ATTN_IMPLEMENTATION = "sdpa" +# The dtype Spyre actually runs in. ``from_pretrained`` defaults to float32, but +# Spyre executes in bfloat16, so both the capture path (``load_model_only``) and +# the YAML emit path (``_tensor_info_to_spec``) default floating-point tensors to +# bfloat16. This keeps the generated config faithful to the runtime dtype +# regardless of the checkpoint's stored precision. Only floating-point dtypes are +# remapped; integer/bool tensors (ids, masks, positions) keep their own dtype. +DEFAULT_FLOAT_DTYPE = torch.bfloat16 +_FLOAT_DTYPE_ALIASES = ("float16", "float32", "float64", "float", "half", "double") + +# Special tensors (position/mask/ids -- see ``_is_special_tensor``) carry indices +# rather than activations, so they are forced to this integer dtype regardless of +# the dtype they were captured under. This makes their ``randint`` init consistent +# (randint on a floating-point dtype is meaningless). +DEFAULT_INT_DTYPE = torch.int64 + def _resolve_attn_implementation(config: Any) -> str: """Return the attention implementation the model actually used. @@ -589,15 +604,49 @@ def _tensor_info_to_spec(tensor_info: Dict[str, Any], name: str) -> Dict[str, An if not dtype.startswith("torch."): dtype = f"torch.{dtype}" + # Default every floating-point tensor to bfloat16 (the dtype Spyre runs in), + # regardless of the precision the checkpoint was captured in. A model loaded + # in float32 would otherwise emit float32 specs; normalizing here guarantees + # the "default is bfloat16" contract even when the capture path did not (or + # could not) load the model in bfloat16. Integer/bool tensors are left alone. + bare_dtype = dtype.replace("torch.", "") + if bare_dtype in _FLOAT_DTYPE_ALIASES: + dtype = str(DEFAULT_FLOAT_DTYPE) + # Determine init strategy based on tensor characteristics is_random = tensor_info.get("is_random", True) init = "randn" if is_random else "zeros" init_args = {} - # Special handling for position/id tensors - if _is_special_tensor(name): + # An integer tensor (e.g. token ids for an embedding, position ids, masks) + # must not use randn -- torch.randn ("normal_kernel_cpu") is float-only and + # raises NotImplementedError for integer dtypes. Use randint for any integer + # dtype, and also for the name-based special tensors (position/mask/ids), + # which may be captured under a generic name like "arg_0". + is_int_dtype = any(t in dtype for t in ("int", "uint", "long", "short", "bool")) + if is_int_dtype or _is_special_tensor(name): init = "randint" - init_args = {"high": 10000} + # A special tensor (position/mask/ids) holds indices, not activations, + # so force it to an integer dtype. This keeps the randint init consistent + # even when the tensor was captured under a floating-point dtype (e.g. a + # "position_embeddings" tensor captured as bfloat16): randint on a float + # dtype is meaningless, so it becomes torch.int64 here. + if _is_special_tensor(name): + dtype = str(DEFAULT_INT_DTYPE) + # Use the smallest dimension of the tensor's own shape as the exclusive + # upper bound (e.g. shape (64, 32, 128) -> high=32). This keeps generated + # index/position values in range for that tensor rather than using a + # fixed, possibly out-of-range constant. Guard against empty shapes and + # zero/one-sized dims (randint needs high >= 1). + shape = tensor_info.get("shape") or [] + high = min(shape) if shape else 1 + init_args = {"high": max(int(high), 1)} + elif init in ("randn", "rand"): + # Float random tensors use xavier init. xavier is undefined for <2-D + # shapes (the OOT framework rejects it), so 1-D float tensors fall back + # to randn. + shape = tensor_info.get("shape") or [] + init = "xavier" if len(shape) >= 2 else "randn" tensor_spec = { "shape": tensor_info["shape"], @@ -940,11 +989,17 @@ def load_model_only( ``Mistral3ForConditionalGeneration`` for VLMs. **from_pretrained_kwargs: Extra kwargs forwarded to ``from_pretrained`` (e.g. ``torch_dtype``, ``device_map``, - ``quantization_config``, ``trust_remote_code``). + ``quantization_config``, ``trust_remote_code``). ``torch_dtype`` + defaults to :data:`DEFAULT_FLOAT_DTYPE` (bfloat16, the dtype Spyre + runs in) rather than ``from_pretrained``'s float32; pass it + explicitly to override. Returns: The loaded, ``.eval()``-mode model. """ + # Capture in bfloat16 by default so the recorded floating-point tensors match + # the dtype Spyre executes in. Callers may still override torch_dtype. + from_pretrained_kwargs.setdefault("torch_dtype", DEFAULT_FLOAT_DTYPE) logger.info(f"Loading model: {model_path} via {model_cls.__name__}") return model_cls.from_pretrained(model_path, **from_pretrained_kwargs).eval() From e52c660b1175336979bd0ddd48776570a6f95a50 Mon Sep 17 00:00:00 2001 From: Kazuaki Ishizaki Date: Wed, 22 Jul 2026 02:32:51 +0000 Subject: [PATCH 2/4] update YAML files Signed-off-by: Kazuaki Ishizaki --- .../Ministral-3-14B-Instruct-2512.yaml | 76 ++++++------ .../Mistral-Small-3.2-24B-Instruct-2506.yaml | 116 +++++++++--------- .../granite_3_3_8b_instruct_spyre.yaml | 68 +++++----- .../module_tests/granite_4_1_8b_spyre.yaml | 104 ++++++++-------- 4 files changed, 186 insertions(+), 178 deletions(-) diff --git a/tests/configs/module_tests/Ministral-3-14B-Instruct-2512.yaml b/tests/configs/module_tests/Ministral-3-14B-Instruct-2512.yaml index 40c9601b..3a2f8a2a 100644 --- a/tests/configs/module_tests/Ministral-3-14B-Instruct-2512.yaml +++ b/tests/configs/module_tests/Ministral-3-14B-Instruct-2512.yaml @@ -46,7 +46,7 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 1 inputs_embeds: tensor: shape: [1, 128, 4096] @@ -54,7 +54,7 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randn + init: xavier - args: [] kwargs: attention_mask: @@ -66,7 +66,7 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 1 inputs_embeds: tensor: shape: [1, 1, 4096] @@ -74,7 +74,7 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randn + init: xavier - name: Ministral3RotaryEmbedding_483e0b9c module_path: transformers.models.ministral3.modeling_ministral3.Ministral3RotaryEmbedding description: 'Module: transformers.models.ministral3.modeling_ministral3.Ministral3RotaryEmbedding' @@ -97,7 +97,7 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: position_ids: tensor: @@ -108,7 +108,7 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 1 - args: - tensor: shape: [1, 1, 4096] @@ -116,7 +116,7 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: position_ids: tensor: @@ -127,7 +127,7 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 1 - name: Ministral3DecoderLayer_layer0 module_path: transformers.models.ministral3.modeling_ministral3.Ministral3DecoderLayer description: 'Module: transformers.models.ministral3.modeling_ministral3.Ministral3DecoderLayer' @@ -151,7 +151,7 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: position_ids: tensor: @@ -162,25 +162,25 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 1 position_embeddings: tensor_list: - shape: [1, 128, 128] stride: null storage_offset: 0 - dtype: torch.bfloat16 + dtype: torch.int64 device: spyre init: randint init_args: - high: 10000 + high: 1 - shape: [1, 128, 128] stride: null storage_offset: 0 - dtype: torch.bfloat16 + dtype: torch.int64 device: spyre init: randint init_args: - high: 10000 + high: 1 - args: - tensor: shape: [1, 1, 4096] @@ -188,7 +188,7 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: position_ids: tensor: @@ -199,25 +199,25 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 1 position_embeddings: tensor_list: - shape: [1, 1, 128] stride: null storage_offset: 0 - dtype: torch.bfloat16 + dtype: torch.int64 device: spyre init: randint init_args: - high: 10000 + high: 1 - shape: [1, 1, 128] stride: null storage_offset: 0 - dtype: torch.bfloat16 + dtype: torch.int64 device: spyre init: randint init_args: - high: 10000 + high: 1 - name: Ministral3RMSNorm_4096 module_path: transformers.models.ministral3.modeling_ministral3.Ministral3RMSNorm description: 'Module: transformers.models.ministral3.modeling_ministral3.Ministral3RMSNorm' @@ -233,7 +233,7 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: {} - args: - tensor: @@ -242,7 +242,7 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: {} - name: Ministral3Attention_layer0 module_path: transformers.models.ministral3.modeling_ministral3.Ministral3Attention @@ -269,7 +269,7 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randn + init: xavier position_ids: tensor: shape: [1, 128] @@ -279,25 +279,25 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 1 position_embeddings: tensor_list: - shape: [1, 128, 128] stride: null storage_offset: 0 - dtype: torch.bfloat16 + dtype: torch.int64 device: spyre init: randint init_args: - high: 10000 + high: 1 - shape: [1, 128, 128] stride: null storage_offset: 0 - dtype: torch.bfloat16 + dtype: torch.int64 device: spyre init: randint init_args: - high: 10000 + high: 1 - args: [] kwargs: hidden_states: @@ -307,7 +307,7 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randn + init: xavier position_ids: tensor: shape: [1, 1] @@ -317,25 +317,25 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 1 position_embeddings: tensor_list: - shape: [1, 1, 128] stride: null storage_offset: 0 - dtype: torch.bfloat16 + dtype: torch.int64 device: spyre init: randint init_args: - high: 10000 + high: 1 - shape: [1, 1, 128] stride: null storage_offset: 0 - dtype: torch.bfloat16 + dtype: torch.int64 device: spyre init: randint init_args: - high: 10000 + high: 1 - name: Ministral3MLP_eb806b7d module_path: transformers.models.ministral3.modeling_ministral3.Ministral3MLP description: 'Module: transformers.models.ministral3.modeling_ministral3.Ministral3MLP' @@ -358,7 +358,7 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: {} - args: - tensor: @@ -367,7 +367,7 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: {} - name: SiLUActivation_d2f532e9 module_path: transformers.activations.SiLUActivation @@ -383,7 +383,7 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: {} - args: - tensor: @@ -392,7 +392,7 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: {} - path: ${TORCH_DEVICE_ROOT}/tests/test_modules_custom.py unlisted_test_mode: skip diff --git a/tests/configs/module_tests/Mistral-Small-3.2-24B-Instruct-2506.yaml b/tests/configs/module_tests/Mistral-Small-3.2-24B-Instruct-2506.yaml index 5b5ba33b..fd3bc747 100644 --- a/tests/configs/module_tests/Mistral-Small-3.2-24B-Instruct-2506.yaml +++ b/tests/configs/module_tests/Mistral-Small-3.2-24B-Instruct-2506.yaml @@ -48,15 +48,15 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 1 inputs_embeds: tensor: shape: [1, 128, 4096] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.bfloat16 device: spyre - init: randn + init: xavier - args: [] kwargs: attention_mask: @@ -68,15 +68,15 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 1 inputs_embeds: tensor: shape: [1, 1, 4096] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.bfloat16 device: spyre - init: randn + init: xavier - name: MistralRotaryEmbedding_0af204ee module_path: transformers.models.mistral.modeling_mistral.MistralRotaryEmbedding description: 'Module: transformers.models.mistral.modeling_mistral.MistralRotaryEmbedding' @@ -97,32 +97,36 @@ test_suite_config: shape: [1, 128, 4096] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.bfloat16 device: spyre - init: randn + init: xavier - tensor: shape: [1, 128] stride: null storage_offset: 0 dtype: torch.int64 device: spyre - init: randn + init: randint + init_args: + high: 1 kwargs: {} - args: - tensor: shape: [1, 1, 4096] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.bfloat16 device: spyre - init: randn + init: xavier - tensor: shape: [1, 1] stride: null storage_offset: 0 dtype: torch.int64 device: spyre - init: randn + init: randint + init_args: + high: 1 kwargs: {} - name: MistralDecoderLayer_layer0 module_path: transformers.models.mistral.modeling_mistral.MistralDecoderLayer @@ -145,9 +149,9 @@ test_suite_config: shape: [1, 128, 4096] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: position_ids: tensor: @@ -158,7 +162,7 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 1 cache_position: tensor: shape: [128] @@ -168,33 +172,33 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 128 position_embeddings: tensor_list: - shape: [1, 128, 128] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.int64 device: spyre init: randint init_args: - high: 10000 + high: 1 - shape: [1, 128, 128] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.int64 device: spyre init: randint init_args: - high: 10000 + high: 1 - args: - tensor: shape: [1, 1, 4096] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: position_ids: tensor: @@ -205,7 +209,7 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 1 cache_position: tensor: shape: [1] @@ -215,25 +219,25 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 1 position_embeddings: tensor_list: - shape: [1, 1, 128] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.int64 device: spyre init: randint init_args: - high: 10000 + high: 1 - shape: [1, 1, 128] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.int64 device: spyre init: randint init_args: - high: 10000 + high: 1 - name: MistralRMSNorm_4096 module_path: transformers.models.mistral.modeling_mistral.MistralRMSNorm description: 'Module: transformers.models.mistral.modeling_mistral.MistralRMSNorm' @@ -247,18 +251,18 @@ test_suite_config: shape: [1, 128, 4096] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: {} - args: - tensor: shape: [1, 1, 4096] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: {} - name: MistralAttention_layer0 module_path: transformers.models.mistral.modeling_mistral.MistralAttention @@ -283,9 +287,9 @@ test_suite_config: shape: [1, 128, 4096] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.bfloat16 device: spyre - init: randn + init: xavier position_ids: tensor: shape: [1, 128] @@ -295,7 +299,7 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 1 cache_position: tensor: shape: [128] @@ -305,25 +309,25 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 128 position_embeddings: tensor_list: - shape: [1, 128, 128] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.int64 device: spyre init: randint init_args: - high: 10000 + high: 1 - shape: [1, 128, 128] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.int64 device: spyre init: randint init_args: - high: 10000 + high: 1 - args: [] kwargs: hidden_states: @@ -331,9 +335,9 @@ test_suite_config: shape: [1, 1, 4096] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.bfloat16 device: spyre - init: randn + init: xavier position_ids: tensor: shape: [1, 1] @@ -343,7 +347,7 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 1 cache_position: tensor: shape: [1] @@ -353,25 +357,25 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 1 position_embeddings: tensor_list: - shape: [1, 1, 128] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.int64 device: spyre init: randint init_args: - high: 10000 + high: 1 - shape: [1, 1, 128] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.int64 device: spyre init: randint init_args: - high: 10000 + high: 1 - name: MistralMLP_f82d9546 module_path: transformers.models.mistral.modeling_mistral.MistralMLP description: 'Module: transformers.models.mistral.modeling_mistral.MistralMLP' @@ -392,18 +396,18 @@ test_suite_config: shape: [1, 128, 4096] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: {} - args: - tensor: shape: [1, 1, 4096] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: {} - name: SiLUActivation_d2f532e9 module_path: transformers.activations.SiLUActivation @@ -417,18 +421,18 @@ test_suite_config: shape: [1, 128, 32768] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: {} - args: - tensor: shape: [1, 1, 32768] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: {} - path: ${TORCH_DEVICE_ROOT}/tests/test_modules_custom.py unlisted_test_mode: skip diff --git a/tests/configs/module_tests/granite_3_3_8b_instruct_spyre.yaml b/tests/configs/module_tests/granite_3_3_8b_instruct_spyre.yaml index ea0d26a7..971b751f 100644 --- a/tests/configs/module_tests/granite_3_3_8b_instruct_spyre.yaml +++ b/tests/configs/module_tests/granite_3_3_8b_instruct_spyre.yaml @@ -38,7 +38,7 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: position_ids: tensor: @@ -49,7 +49,7 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 1 - args: - tensor: shape: [1, 1, 4096] @@ -57,7 +57,7 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: position_ids: tensor: @@ -68,7 +68,7 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 1 - name: GraniteDecoderLayer_layer0 module_path: transformers.models.granite.modeling_granite.GraniteDecoderLayer description: 'Module: transformers.models.granite.modeling_granite.GraniteDecoderLayer' @@ -92,7 +92,7 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: position_ids: tensor: @@ -103,25 +103,25 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 1 position_embeddings: tensor_list: - shape: [1, 128, 128] stride: null storage_offset: 0 - dtype: torch.bfloat16 + dtype: torch.int64 device: spyre init: randint init_args: - high: 10000 + high: 1 - shape: [1, 128, 128] stride: null storage_offset: 0 - dtype: torch.bfloat16 + dtype: torch.int64 device: spyre init: randint init_args: - high: 10000 + high: 1 - args: - tensor: shape: [1, 1, 4096] @@ -129,7 +129,7 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: position_ids: tensor: @@ -140,25 +140,25 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 1 position_embeddings: tensor_list: - shape: [1, 1, 128] stride: null storage_offset: 0 - dtype: torch.bfloat16 + dtype: torch.int64 device: spyre init: randint init_args: - high: 10000 + high: 1 - shape: [1, 1, 128] stride: null storage_offset: 0 - dtype: torch.bfloat16 + dtype: torch.int64 device: spyre init: randint init_args: - high: 10000 + high: 1 - name: GraniteRMSNorm_4096 module_path: transformers.models.granite.modeling_granite.GraniteRMSNorm description: 'Module: transformers.models.granite.modeling_granite.GraniteRMSNorm' @@ -174,7 +174,7 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: {} - args: - tensor: @@ -183,7 +183,7 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: {} - name: GraniteAttention_layer0 module_path: transformers.models.granite.modeling_granite.GraniteAttention @@ -210,7 +210,7 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randn + init: xavier position_ids: tensor: shape: [1, 128] @@ -220,25 +220,25 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 1 position_embeddings: tensor_list: - shape: [1, 128, 128] stride: null storage_offset: 0 - dtype: torch.bfloat16 + dtype: torch.int64 device: spyre init: randint init_args: - high: 10000 + high: 1 - shape: [1, 128, 128] stride: null storage_offset: 0 - dtype: torch.bfloat16 + dtype: torch.int64 device: spyre init: randint init_args: - high: 10000 + high: 1 - args: [] kwargs: hidden_states: @@ -248,7 +248,7 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randn + init: xavier position_ids: tensor: shape: [1, 1] @@ -258,25 +258,25 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 1 position_embeddings: tensor_list: - shape: [1, 1, 128] stride: null storage_offset: 0 - dtype: torch.bfloat16 + dtype: torch.int64 device: spyre init: randint init_args: - high: 10000 + high: 1 - shape: [1, 1, 128] stride: null storage_offset: 0 - dtype: torch.bfloat16 + dtype: torch.int64 device: spyre init: randint init_args: - high: 10000 + high: 1 - name: GraniteMLP_77d2613c module_path: transformers.models.granite.modeling_granite.GraniteMLP description: 'Module: transformers.models.granite.modeling_granite.GraniteMLP' @@ -299,7 +299,7 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: {} - args: - tensor: @@ -308,7 +308,7 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: {} - name: SiLUActivation_d2f532e9 module_path: transformers.activations.SiLUActivation @@ -324,7 +324,7 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: {} - args: - tensor: @@ -333,7 +333,7 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: {} - path: ${TORCH_DEVICE_ROOT}/tests/test_modules_custom.py unlisted_test_mode: skip diff --git a/tests/configs/module_tests/granite_4_1_8b_spyre.yaml b/tests/configs/module_tests/granite_4_1_8b_spyre.yaml index 79e1aa56..ee2e855f 100644 --- a/tests/configs/module_tests/granite_4_1_8b_spyre.yaml +++ b/tests/configs/module_tests/granite_4_1_8b_spyre.yaml @@ -36,32 +36,36 @@ test_suite_config: shape: [1, 128, 4096] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.bfloat16 device: spyre - init: randn + init: xavier - tensor: shape: [1, 128] stride: null storage_offset: 0 dtype: torch.int64 device: spyre - init: randn + init: randint + init_args: + high: 1 kwargs: {} - args: - tensor: shape: [1, 1, 4096] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.bfloat16 device: spyre - init: randn + init: xavier - tensor: shape: [1, 1] stride: null storage_offset: 0 dtype: torch.int64 device: spyre - init: randn + init: randint + init_args: + high: 1 kwargs: {} - name: GraniteDecoderLayer_layer0 module_path: transformers.models.granite.modeling_granite.GraniteDecoderLayer @@ -84,9 +88,9 @@ test_suite_config: shape: [1, 128, 4096] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: position_ids: tensor: @@ -97,7 +101,7 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 1 cache_position: tensor: shape: [128] @@ -107,33 +111,33 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 128 position_embeddings: tensor_list: - shape: [1, 128, 128] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.int64 device: spyre init: randint init_args: - high: 10000 + high: 1 - shape: [1, 128, 128] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.int64 device: spyre init: randint init_args: - high: 10000 + high: 1 - args: - tensor: shape: [1, 1, 4096] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: position_ids: tensor: @@ -144,7 +148,7 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 1 cache_position: tensor: shape: [1] @@ -154,25 +158,25 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 1 position_embeddings: tensor_list: - shape: [1, 1, 128] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.int64 device: spyre init: randint init_args: - high: 10000 + high: 1 - shape: [1, 1, 128] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.int64 device: spyre init: randint init_args: - high: 10000 + high: 1 - name: GraniteRMSNorm_4096 module_path: transformers.models.granite.modeling_granite.GraniteRMSNorm description: 'Module: transformers.models.granite.modeling_granite.GraniteRMSNorm' @@ -186,18 +190,18 @@ test_suite_config: shape: [1, 128, 4096] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: {} - args: - tensor: shape: [1, 1, 4096] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: {} - name: GraniteAttention_layer0 module_path: transformers.models.granite.modeling_granite.GraniteAttention @@ -222,9 +226,9 @@ test_suite_config: shape: [1, 128, 4096] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.bfloat16 device: spyre - init: randn + init: xavier position_ids: tensor: shape: [1, 128] @@ -234,7 +238,7 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 1 cache_position: tensor: shape: [128] @@ -244,25 +248,25 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 128 position_embeddings: tensor_list: - shape: [1, 128, 128] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.int64 device: spyre init: randint init_args: - high: 10000 + high: 1 - shape: [1, 128, 128] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.int64 device: spyre init: randint init_args: - high: 10000 + high: 1 - args: [] kwargs: hidden_states: @@ -270,9 +274,9 @@ test_suite_config: shape: [1, 1, 4096] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.bfloat16 device: spyre - init: randn + init: xavier position_ids: tensor: shape: [1, 1] @@ -282,7 +286,7 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 1 cache_position: tensor: shape: [1] @@ -292,25 +296,25 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 1 position_embeddings: tensor_list: - shape: [1, 1, 128] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.int64 device: spyre init: randint init_args: - high: 10000 + high: 1 - shape: [1, 1, 128] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.int64 device: spyre init: randint init_args: - high: 10000 + high: 1 - name: GraniteMLP_77d2613c module_path: transformers.models.granite.modeling_granite.GraniteMLP description: 'Module: transformers.models.granite.modeling_granite.GraniteMLP' @@ -331,18 +335,18 @@ test_suite_config: shape: [1, 128, 4096] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: {} - args: - tensor: shape: [1, 1, 4096] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: {} - name: SiLUActivation_d2f532e9 module_path: transformers.activations.SiLUActivation @@ -356,18 +360,18 @@ test_suite_config: shape: [1, 128, 12800] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: {} - args: - tensor: shape: [1, 1, 12800] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: {} - path: ${TORCH_DEVICE_ROOT}/tests/test_modules_custom.py unlisted_test_mode: skip From e87024780d96e8a2406508b9f77042810f64d0a9 Mon Sep 17 00:00:00 2001 From: Kazuaki Ishizaki Date: Wed, 22 Jul 2026 03:04:27 +0000 Subject: [PATCH 3/4] Fix handling of position_embeddings Signed-off-by: Kazuaki Ishizaki --- .../Ministral-3-14B-Instruct-2512.yaml | 48 +++++++------------ .../Mistral-Small-3.2-24B-Instruct-2506.yaml | 48 +++++++------------ .../granite_3_3_8b_instruct_spyre.yaml | 48 +++++++------------ .../module_tests/granite_4_1_8b_spyre.yaml | 48 +++++++------------ .../auto_generate_module_config.py | 4 +- 5 files changed, 67 insertions(+), 129 deletions(-) diff --git a/tests/configs/module_tests/Ministral-3-14B-Instruct-2512.yaml b/tests/configs/module_tests/Ministral-3-14B-Instruct-2512.yaml index 3a2f8a2a..f0e5c91e 100644 --- a/tests/configs/module_tests/Ministral-3-14B-Instruct-2512.yaml +++ b/tests/configs/module_tests/Ministral-3-14B-Instruct-2512.yaml @@ -168,19 +168,15 @@ test_suite_config: - shape: [1, 128, 128] stride: null storage_offset: 0 - dtype: torch.int64 + dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 1 + init: xavier - shape: [1, 128, 128] stride: null storage_offset: 0 - dtype: torch.int64 + dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 1 + init: xavier - args: - tensor: shape: [1, 1, 4096] @@ -205,19 +201,15 @@ test_suite_config: - shape: [1, 1, 128] stride: null storage_offset: 0 - dtype: torch.int64 + dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 1 + init: xavier - shape: [1, 1, 128] stride: null storage_offset: 0 - dtype: torch.int64 + dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 1 + init: xavier - name: Ministral3RMSNorm_4096 module_path: transformers.models.ministral3.modeling_ministral3.Ministral3RMSNorm description: 'Module: transformers.models.ministral3.modeling_ministral3.Ministral3RMSNorm' @@ -285,19 +277,15 @@ test_suite_config: - shape: [1, 128, 128] stride: null storage_offset: 0 - dtype: torch.int64 + dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 1 + init: xavier - shape: [1, 128, 128] stride: null storage_offset: 0 - dtype: torch.int64 + dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 1 + init: xavier - args: [] kwargs: hidden_states: @@ -323,19 +311,15 @@ test_suite_config: - shape: [1, 1, 128] stride: null storage_offset: 0 - dtype: torch.int64 + dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 1 + init: xavier - shape: [1, 1, 128] stride: null storage_offset: 0 - dtype: torch.int64 + dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 1 + init: xavier - name: Ministral3MLP_eb806b7d module_path: transformers.models.ministral3.modeling_ministral3.Ministral3MLP description: 'Module: transformers.models.ministral3.modeling_ministral3.Ministral3MLP' diff --git a/tests/configs/module_tests/Mistral-Small-3.2-24B-Instruct-2506.yaml b/tests/configs/module_tests/Mistral-Small-3.2-24B-Instruct-2506.yaml index fd3bc747..c0add4f8 100644 --- a/tests/configs/module_tests/Mistral-Small-3.2-24B-Instruct-2506.yaml +++ b/tests/configs/module_tests/Mistral-Small-3.2-24B-Instruct-2506.yaml @@ -178,19 +178,15 @@ test_suite_config: - shape: [1, 128, 128] stride: null storage_offset: 0 - dtype: torch.int64 + dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 1 + init: xavier - shape: [1, 128, 128] stride: null storage_offset: 0 - dtype: torch.int64 + dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 1 + init: xavier - args: - tensor: shape: [1, 1, 4096] @@ -225,19 +221,15 @@ test_suite_config: - shape: [1, 1, 128] stride: null storage_offset: 0 - dtype: torch.int64 + dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 1 + init: xavier - shape: [1, 1, 128] stride: null storage_offset: 0 - dtype: torch.int64 + dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 1 + init: xavier - name: MistralRMSNorm_4096 module_path: transformers.models.mistral.modeling_mistral.MistralRMSNorm description: 'Module: transformers.models.mistral.modeling_mistral.MistralRMSNorm' @@ -315,19 +307,15 @@ test_suite_config: - shape: [1, 128, 128] stride: null storage_offset: 0 - dtype: torch.int64 + dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 1 + init: xavier - shape: [1, 128, 128] stride: null storage_offset: 0 - dtype: torch.int64 + dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 1 + init: xavier - args: [] kwargs: hidden_states: @@ -363,19 +351,15 @@ test_suite_config: - shape: [1, 1, 128] stride: null storage_offset: 0 - dtype: torch.int64 + dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 1 + init: xavier - shape: [1, 1, 128] stride: null storage_offset: 0 - dtype: torch.int64 + dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 1 + init: xavier - name: MistralMLP_f82d9546 module_path: transformers.models.mistral.modeling_mistral.MistralMLP description: 'Module: transformers.models.mistral.modeling_mistral.MistralMLP' diff --git a/tests/configs/module_tests/granite_3_3_8b_instruct_spyre.yaml b/tests/configs/module_tests/granite_3_3_8b_instruct_spyre.yaml index 971b751f..c782b4e6 100644 --- a/tests/configs/module_tests/granite_3_3_8b_instruct_spyre.yaml +++ b/tests/configs/module_tests/granite_3_3_8b_instruct_spyre.yaml @@ -109,19 +109,15 @@ test_suite_config: - shape: [1, 128, 128] stride: null storage_offset: 0 - dtype: torch.int64 + dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 1 + init: xavier - shape: [1, 128, 128] stride: null storage_offset: 0 - dtype: torch.int64 + dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 1 + init: xavier - args: - tensor: shape: [1, 1, 4096] @@ -146,19 +142,15 @@ test_suite_config: - shape: [1, 1, 128] stride: null storage_offset: 0 - dtype: torch.int64 + dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 1 + init: xavier - shape: [1, 1, 128] stride: null storage_offset: 0 - dtype: torch.int64 + dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 1 + init: xavier - name: GraniteRMSNorm_4096 module_path: transformers.models.granite.modeling_granite.GraniteRMSNorm description: 'Module: transformers.models.granite.modeling_granite.GraniteRMSNorm' @@ -226,19 +218,15 @@ test_suite_config: - shape: [1, 128, 128] stride: null storage_offset: 0 - dtype: torch.int64 + dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 1 + init: xavier - shape: [1, 128, 128] stride: null storage_offset: 0 - dtype: torch.int64 + dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 1 + init: xavier - args: [] kwargs: hidden_states: @@ -264,19 +252,15 @@ test_suite_config: - shape: [1, 1, 128] stride: null storage_offset: 0 - dtype: torch.int64 + dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 1 + init: xavier - shape: [1, 1, 128] stride: null storage_offset: 0 - dtype: torch.int64 + dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 1 + init: xavier - name: GraniteMLP_77d2613c module_path: transformers.models.granite.modeling_granite.GraniteMLP description: 'Module: transformers.models.granite.modeling_granite.GraniteMLP' diff --git a/tests/configs/module_tests/granite_4_1_8b_spyre.yaml b/tests/configs/module_tests/granite_4_1_8b_spyre.yaml index ee2e855f..1b0ce3f1 100644 --- a/tests/configs/module_tests/granite_4_1_8b_spyre.yaml +++ b/tests/configs/module_tests/granite_4_1_8b_spyre.yaml @@ -117,19 +117,15 @@ test_suite_config: - shape: [1, 128, 128] stride: null storage_offset: 0 - dtype: torch.int64 + dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 1 + init: xavier - shape: [1, 128, 128] stride: null storage_offset: 0 - dtype: torch.int64 + dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 1 + init: xavier - args: - tensor: shape: [1, 1, 4096] @@ -164,19 +160,15 @@ test_suite_config: - shape: [1, 1, 128] stride: null storage_offset: 0 - dtype: torch.int64 + dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 1 + init: xavier - shape: [1, 1, 128] stride: null storage_offset: 0 - dtype: torch.int64 + dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 1 + init: xavier - name: GraniteRMSNorm_4096 module_path: transformers.models.granite.modeling_granite.GraniteRMSNorm description: 'Module: transformers.models.granite.modeling_granite.GraniteRMSNorm' @@ -254,19 +246,15 @@ test_suite_config: - shape: [1, 128, 128] stride: null storage_offset: 0 - dtype: torch.int64 + dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 1 + init: xavier - shape: [1, 128, 128] stride: null storage_offset: 0 - dtype: torch.int64 + dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 1 + init: xavier - args: [] kwargs: hidden_states: @@ -302,19 +290,15 @@ test_suite_config: - shape: [1, 1, 128] stride: null storage_offset: 0 - dtype: torch.int64 + dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 1 + init: xavier - shape: [1, 1, 128] stride: null storage_offset: 0 - dtype: torch.int64 + dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 1 + init: xavier - name: GraniteMLP_77d2613c module_path: transformers.models.granite.modeling_granite.GraniteMLP description: 'Module: transformers.models.granite.modeling_granite.GraniteMLP' diff --git a/utils/module_discovery/auto_generate_module_config.py b/utils/module_discovery/auto_generate_module_config.py index d7add9e5..95b20d84 100644 --- a/utils/module_discovery/auto_generate_module_config.py +++ b/utils/module_discovery/auto_generate_module_config.py @@ -72,7 +72,9 @@ def represent_data(self, data): def _is_special_tensor(name: str) -> bool: """Check if tensor name indicates it should not be random.""" - return any(keyword in name.lower() for keyword in ["position", "mask", "ids"]) + return "position_embedding" not in name.lower() and any( + keyword in name.lower() for keyword in ["position", "mask", "ids"] + ) # Extracted from the loaded config so a standalone module rebuilt from the YAML From ab353de05e85a147d55e60843c822b62cdba908f Mon Sep 17 00:00:00 2001 From: Kazuaki Ishizaki Date: Wed, 22 Jul 2026 03:13:16 +0000 Subject: [PATCH 4/4] Fix comments Signed-off-by: Kazuaki Ishizaki --- utils/module_discovery/auto_generate_module_config.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/utils/module_discovery/auto_generate_module_config.py b/utils/module_discovery/auto_generate_module_config.py index 95b20d84..481b36b8 100644 --- a/utils/module_discovery/auto_generate_module_config.py +++ b/utils/module_discovery/auto_generate_module_config.py @@ -630,8 +630,8 @@ def _tensor_info_to_spec(tensor_info: Dict[str, Any], name: str) -> Dict[str, An init = "randint" # A special tensor (position/mask/ids) holds indices, not activations, # so force it to an integer dtype. This keeps the randint init consistent - # even when the tensor was captured under a floating-point dtype (e.g. a - # "position_embeddings" tensor captured as bfloat16): randint on a float + # even when the tensor was captured under a floating-point dtype + # tensor captured as bfloat16): randint on a float # dtype is meaningless, so it becomes torch.int64 here. if _is_special_tensor(name): dtype = str(DEFAULT_INT_DTYPE)