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..f0e5c91e 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,7 +162,7 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 1 position_embeddings: tensor_list: - shape: [1, 128, 128] @@ -170,17 +170,13 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 10000 + init: xavier - shape: [1, 128, 128] stride: null storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 10000 + init: xavier - args: - tensor: shape: [1, 1, 4096] @@ -188,7 +184,7 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: position_ids: tensor: @@ -199,7 +195,7 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 1 position_embeddings: tensor_list: - shape: [1, 1, 128] @@ -207,17 +203,13 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 10000 + init: xavier - shape: [1, 1, 128] stride: null storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 10000 + init: xavier - name: Ministral3RMSNorm_4096 module_path: transformers.models.ministral3.modeling_ministral3.Ministral3RMSNorm description: 'Module: transformers.models.ministral3.modeling_ministral3.Ministral3RMSNorm' @@ -233,7 +225,7 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: {} - args: - tensor: @@ -242,7 +234,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 +261,7 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randn + init: xavier position_ids: tensor: shape: [1, 128] @@ -279,7 +271,7 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 1 position_embeddings: tensor_list: - shape: [1, 128, 128] @@ -287,17 +279,13 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 10000 + init: xavier - shape: [1, 128, 128] stride: null storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 10000 + init: xavier - args: [] kwargs: hidden_states: @@ -307,7 +295,7 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randn + init: xavier position_ids: tensor: shape: [1, 1] @@ -317,7 +305,7 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 1 position_embeddings: tensor_list: - shape: [1, 1, 128] @@ -325,17 +313,13 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 10000 + init: xavier - shape: [1, 1, 128] stride: null storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 10000 + init: xavier - name: Ministral3MLP_eb806b7d module_path: transformers.models.ministral3.modeling_ministral3.Ministral3MLP description: 'Module: transformers.models.ministral3.modeling_ministral3.Ministral3MLP' @@ -358,7 +342,7 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: {} - args: - tensor: @@ -367,7 +351,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 +367,7 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: {} - args: - tensor: @@ -392,7 +376,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..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 @@ -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,29 @@ 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.bfloat16 device: spyre - init: randint - init_args: - high: 10000 + init: xavier - shape: [1, 128, 128] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 10000 + init: xavier - 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 +205,7 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 1 cache_position: tensor: shape: [1] @@ -215,25 +215,21 @@ 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.bfloat16 device: spyre - init: randint - init_args: - high: 10000 + init: xavier - shape: [1, 1, 128] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 10000 + init: xavier - name: MistralRMSNorm_4096 module_path: transformers.models.mistral.modeling_mistral.MistralRMSNorm description: 'Module: transformers.models.mistral.modeling_mistral.MistralRMSNorm' @@ -247,18 +243,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 +279,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 +291,7 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 1 cache_position: tensor: shape: [128] @@ -305,25 +301,21 @@ 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.bfloat16 device: spyre - init: randint - init_args: - high: 10000 + init: xavier - shape: [1, 128, 128] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 10000 + init: xavier - args: [] kwargs: hidden_states: @@ -331,9 +323,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 +335,7 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 1 cache_position: tensor: shape: [1] @@ -353,25 +345,21 @@ 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.bfloat16 device: spyre - init: randint - init_args: - high: 10000 + init: xavier - shape: [1, 1, 128] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 10000 + init: xavier - name: MistralMLP_f82d9546 module_path: transformers.models.mistral.modeling_mistral.MistralMLP description: 'Module: transformers.models.mistral.modeling_mistral.MistralMLP' @@ -392,18 +380,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 +405,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..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 @@ -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,7 +103,7 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 1 position_embeddings: tensor_list: - shape: [1, 128, 128] @@ -111,17 +111,13 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 10000 + init: xavier - shape: [1, 128, 128] stride: null storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 10000 + init: xavier - args: - tensor: shape: [1, 1, 4096] @@ -129,7 +125,7 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: position_ids: tensor: @@ -140,7 +136,7 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 1 position_embeddings: tensor_list: - shape: [1, 1, 128] @@ -148,17 +144,13 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 10000 + init: xavier - shape: [1, 1, 128] stride: null storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 10000 + init: xavier - name: GraniteRMSNorm_4096 module_path: transformers.models.granite.modeling_granite.GraniteRMSNorm description: 'Module: transformers.models.granite.modeling_granite.GraniteRMSNorm' @@ -174,7 +166,7 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: {} - args: - tensor: @@ -183,7 +175,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 +202,7 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randn + init: xavier position_ids: tensor: shape: [1, 128] @@ -220,7 +212,7 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 1 position_embeddings: tensor_list: - shape: [1, 128, 128] @@ -228,17 +220,13 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 10000 + init: xavier - shape: [1, 128, 128] stride: null storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 10000 + init: xavier - args: [] kwargs: hidden_states: @@ -248,7 +236,7 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randn + init: xavier position_ids: tensor: shape: [1, 1] @@ -258,7 +246,7 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 1 position_embeddings: tensor_list: - shape: [1, 1, 128] @@ -266,17 +254,13 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 10000 + init: xavier - shape: [1, 1, 128] stride: null storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 10000 + init: xavier - name: GraniteMLP_77d2613c module_path: transformers.models.granite.modeling_granite.GraniteMLP description: 'Module: transformers.models.granite.modeling_granite.GraniteMLP' @@ -299,7 +283,7 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: {} - args: - tensor: @@ -308,7 +292,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 +308,7 @@ test_suite_config: storage_offset: 0 dtype: torch.bfloat16 device: spyre - init: randn + init: xavier kwargs: {} - args: - tensor: @@ -333,7 +317,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..1b0ce3f1 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,29 @@ 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.bfloat16 device: spyre - init: randint - init_args: - high: 10000 + init: xavier - shape: [1, 128, 128] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 10000 + init: xavier - 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 +144,7 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 1 cache_position: tensor: shape: [1] @@ -154,25 +154,21 @@ 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.bfloat16 device: spyre - init: randint - init_args: - high: 10000 + init: xavier - shape: [1, 1, 128] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 10000 + init: xavier - name: GraniteRMSNorm_4096 module_path: transformers.models.granite.modeling_granite.GraniteRMSNorm description: 'Module: transformers.models.granite.modeling_granite.GraniteRMSNorm' @@ -186,18 +182,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 +218,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 +230,7 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 1 cache_position: tensor: shape: [128] @@ -244,25 +240,21 @@ 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.bfloat16 device: spyre - init: randint - init_args: - high: 10000 + init: xavier - shape: [1, 128, 128] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 10000 + init: xavier - args: [] kwargs: hidden_states: @@ -270,9 +262,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 +274,7 @@ test_suite_config: device: spyre init: randint init_args: - high: 10000 + high: 1 cache_position: tensor: shape: [1] @@ -292,25 +284,21 @@ 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.bfloat16 device: spyre - init: randint - init_args: - high: 10000 + init: xavier - shape: [1, 1, 128] stride: null storage_offset: 0 - dtype: torch.float32 + dtype: torch.bfloat16 device: spyre - init: randint - init_args: - high: 10000 + init: xavier - name: GraniteMLP_77d2613c module_path: transformers.models.granite.modeling_granite.GraniteMLP description: 'Module: transformers.models.granite.modeling_granite.GraniteMLP' @@ -331,18 +319,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 +344,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 diff --git a/utils/module_discovery/auto_generate_module_config.py b/utils/module_discovery/auto_generate_module_config.py index e32182ce..481b36b8 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 @@ -83,6 +85,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,6 +606,15 @@ 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" @@ -602,6 +628,13 @@ def _tensor_info_to_spec(tensor_info: Dict[str, Any], name: str) -> Dict[str, An 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" + # 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 + # 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 @@ -958,11 +991,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()