| { |
| "architecture": { |
| "attention_variant": "MHA", |
| "family": "masked_lm", |
| "is_moe": false, |
| "mixer": "attention", |
| "positional": "learned", |
| "tie_word_embeddings": true, |
| "view": "encoder" |
| }, |
| "capabilities": { |
| "attention_backends": [ |
| "eager", |
| "sdpa", |
| "flash_attention", |
| "flex_attention" |
| ], |
| "attention_patterns": [ |
| "bidirectional" |
| ], |
| "attention_schedule": null, |
| "task_heads": [ |
| "masked_lm", |
| "multiple_choice", |
| "pretraining", |
| "question_answering", |
| "sequence_classification", |
| "text_encoding", |
| "token_classification" |
| ], |
| "tensor_parallel": false |
| }, |
| "components": [ |
| { |
| "children": [ |
| "embeddings", |
| "encoder", |
| "pooler", |
| "pooler_activation" |
| ], |
| "class_name": "AlbertModel", |
| "id": "model", |
| "kind": "model", |
| "path_pattern": "model" |
| }, |
| { |
| "attributes": { |
| "embedding_dim": "config.hidden_size", |
| "num_embeddings": "config.vocab_size" |
| }, |
| "children": [ |
| "embeddings.word_embeddings", |
| "embeddings.position_embeddings", |
| "embeddings.token_type_embeddings", |
| "embeddings.LayerNorm" |
| ], |
| "class_name": "AlbertEmbeddings", |
| "id": "embeddings", |
| "kind": "embedding", |
| "path_pattern": "model.embeddings" |
| }, |
| { |
| "attributes": { |
| "embedding_dim": 128, |
| "num_embeddings": "config.vocab_size", |
| "tied_lm_head": true |
| }, |
| "class_name": "Embedding", |
| "id": "embeddings.word_embeddings", |
| "kind": "embedding", |
| "path_pattern": "model.embeddings.word_embeddings" |
| }, |
| { |
| "attributes": { |
| "max_position_embeddings": "config.max_position_embeddings", |
| "scheme": "learned" |
| }, |
| "class_name": "Embedding", |
| "id": "embeddings.position_embeddings", |
| "kind": "position", |
| "path_pattern": "model.embeddings.position_embeddings" |
| }, |
| { |
| "attributes": { |
| "embedding_dim": 128, |
| "num_embeddings": 2 |
| }, |
| "class_name": "Embedding", |
| "id": "embeddings.token_type_embeddings", |
| "kind": "embedding", |
| "path_pattern": "model.embeddings.token_type_embeddings" |
| }, |
| { |
| "attributes": { |
| "norm_type": "layer" |
| }, |
| "class_name": "LayerNorm", |
| "id": "embeddings.LayerNorm", |
| "kind": "normalization", |
| "path_pattern": "model.embeddings.LayerNorm" |
| }, |
| { |
| "children": [ |
| "encoder.albert_layer_groups" |
| ], |
| "class_name": "AlbertTransformer", |
| "id": "encoder", |
| "kind": "encoder", |
| "path_pattern": "model.encoder" |
| }, |
| { |
| "children": [ |
| "encoder.albert_layer_groups.0" |
| ], |
| "class_name": "ModuleList", |
| "id": "encoder.albert_layer_groups", |
| "kind": "repeated_container", |
| "path_pattern": "model.encoder.albert_layer_groups" |
| }, |
| { |
| "children": [ |
| "encoder.albert_layer_groups.0.albert_layers" |
| ], |
| "class_name": "AlbertLayerGroup", |
| "id": "encoder.albert_layer_groups.0", |
| "kind": "module", |
| "path_pattern": "model.encoder.albert_layer_groups.0" |
| }, |
| { |
| "children": [ |
| "encoder.albert_layer_groups.0.albert_layers.0" |
| ], |
| "class_name": "ModuleList", |
| "id": "encoder.albert_layer_groups.0.albert_layers", |
| "kind": "repeated_container", |
| "path_pattern": "model.encoder.albert_layer_groups.0.albert_layers" |
| }, |
| { |
| "children": [ |
| "encoder.albert_layer_groups.0.albert_layers.0.full_layer_layer_norm", |
| "encoder.albert_layer_groups.0.albert_layers.0.attention" |
| ], |
| "class_name": "AlbertLayer", |
| "id": "encoder.albert_layer_groups.0.albert_layers.0", |
| "kind": "transformer_block", |
| "path_pattern": "model.encoder.albert_layer_groups.0.albert_layers.0" |
| }, |
| { |
| "attributes": { |
| "norm_type": "layer" |
| }, |
| "class_name": "LayerNorm", |
| "id": "encoder.albert_layer_groups.0.albert_layers.0.full_layer_layer_norm", |
| "kind": "normalization", |
| "path_pattern": "model.encoder.albert_layer_groups.0.albert_layers.0.full_layer_layer_norm" |
| }, |
| { |
| "attributes": { |
| "head_dim": 64, |
| "n_heads": 64, |
| "n_kv_heads": 64, |
| "pattern": "bidirectional", |
| "rope": false, |
| "variant": "MHA" |
| }, |
| "children": [ |
| "encoder.albert_layer_groups.0.albert_layers.0.attention.query", |
| "encoder.albert_layer_groups.0.albert_layers.0.attention.key", |
| "encoder.albert_layer_groups.0.albert_layers.0.attention.value", |
| "encoder.albert_layer_groups.0.albert_layers.0.attention.dense", |
| "encoder.albert_layer_groups.0.albert_layers.0.attention.LayerNorm" |
| ], |
| "class_name": "AlbertAttention", |
| "id": "encoder.albert_layer_groups.0.albert_layers.0.attention", |
| "kind": "attention", |
| "path_pattern": "model.encoder.albert_layer_groups.0.albert_layers.0.attention" |
| }, |
| { |
| "attributes": { |
| "in_features": "config.hidden_size", |
| "out_features": "config.hidden_size" |
| }, |
| "class_name": "Linear", |
| "id": "encoder.albert_layer_groups.0.albert_layers.0.attention.query", |
| "kind": "projection", |
| "path_pattern": "model.encoder.albert_layer_groups.0.albert_layers.0.attention.query" |
| }, |
| { |
| "attributes": { |
| "in_features": "config.hidden_size", |
| "out_features": "config.hidden_size" |
| }, |
| "class_name": "Linear", |
| "id": "encoder.albert_layer_groups.0.albert_layers.0.attention.key", |
| "kind": "projection", |
| "path_pattern": "model.encoder.albert_layer_groups.0.albert_layers.0.attention.key" |
| }, |
| { |
| "attributes": { |
| "in_features": "config.hidden_size", |
| "out_features": "config.hidden_size" |
| }, |
| "class_name": "Linear", |
| "id": "encoder.albert_layer_groups.0.albert_layers.0.attention.value", |
| "kind": "projection", |
| "path_pattern": "model.encoder.albert_layer_groups.0.albert_layers.0.attention.value" |
| }, |
| { |
| "attributes": { |
| "in_features": "config.hidden_size", |
| "out_features": "config.hidden_size" |
| }, |
| "class_name": "Linear", |
| "id": "encoder.albert_layer_groups.0.albert_layers.0.attention.dense", |
| "kind": "projection", |
| "path_pattern": "model.encoder.albert_layer_groups.0.albert_layers.0.attention.dense" |
| }, |
| { |
| "attributes": { |
| "norm_type": "layer" |
| }, |
| "class_name": "LayerNorm", |
| "id": "encoder.albert_layer_groups.0.albert_layers.0.attention.LayerNorm", |
| "kind": "normalization", |
| "path_pattern": "model.encoder.albert_layer_groups.0.albert_layers.0.attention.LayerNorm" |
| }, |
| { |
| "attributes": { |
| "in_features": "config.hidden_size", |
| "out_features": "config.hidden_size" |
| }, |
| "class_name": "Linear", |
| "id": "pooler", |
| "kind": "projection", |
| "path_pattern": "model.pooler" |
| }, |
| { |
| "class_name": "Tanh", |
| "id": "pooler_activation", |
| "kind": "module", |
| "path_pattern": "model.pooler_activation" |
| } |
| ], |
| "config": { |
| "class_name": "AlbertConfig", |
| "model_type": "albert", |
| "module": "transformers.models.albert.configuration_albert", |
| "referenced_fields": {}, |
| "salient_fields": { |
| "hidden_act": "gelu_new", |
| "hidden_size": 4096, |
| "intermediate_size": 16384, |
| "is_encoder_decoder": false, |
| "max_position_embeddings": 512, |
| "num_attention_heads": 64, |
| "num_hidden_layers": 12, |
| "tie_word_embeddings": true, |
| "vocab_size": 30000 |
| } |
| }, |
| "dataflow": { |
| "input": { |
| "name": "input_ids", |
| "shape": [ |
| "B", |
| "S" |
| ] |
| }, |
| "output": { |
| "shape": [ |
| "B", |
| "S", |
| "config.hidden_size" |
| ] |
| }, |
| "shapes": { |
| "embeddings": { |
| "in": [ |
| "B", |
| "S" |
| ], |
| "out": [ |
| "B", |
| "S", |
| 128 |
| ] |
| }, |
| "encoder": { |
| "in": [ |
| "B", |
| "S", |
| 128 |
| ], |
| "out": [ |
| "B", |
| "S", |
| "config.hidden_size" |
| ] |
| }, |
| "pooler": { |
| "in": [ |
| "B", |
| "config.hidden_size" |
| ], |
| "out": [ |
| "B", |
| "config.hidden_size" |
| ] |
| }, |
| "pooler_activation": { |
| "in": [ |
| "B", |
| "config.hidden_size" |
| ], |
| "out": [ |
| "B", |
| "config.hidden_size" |
| ] |
| } |
| }, |
| "source": "observed_forward_meta" |
| }, |
| "edges": [ |
| { |
| "kind": "data", |
| "source": "embeddings", |
| "target": "encoder" |
| }, |
| { |
| "kind": "data", |
| "source": "embeddings.word_embeddings", |
| "target": "embeddings.token_type_embeddings" |
| }, |
| { |
| "kind": "data", |
| "source": "embeddings.token_type_embeddings", |
| "target": "embeddings.LayerNorm" |
| }, |
| { |
| "kind": "data", |
| "source": "encoder.albert_layer_groups.0.albert_layers.0.full_layer_layer_norm", |
| "target": "encoder.albert_layer_groups.0.albert_layers.0.attention" |
| }, |
| { |
| "kind": "mask", |
| "source": "input:attention_mask", |
| "target": "encoder.albert_layer_groups.0.albert_layers.0.attention" |
| }, |
| { |
| "kind": "residual", |
| "source": "encoder.albert_layer_groups.0.albert_layers.0", |
| "target": "encoder.albert_layer_groups.0.albert_layers.0.attention" |
| }, |
| { |
| "kind": "position", |
| "source": "embeddings.position_embeddings", |
| "target": "embeddings" |
| }, |
| { |
| "kind": "data", |
| "provenance": "intra_module", |
| "source": "encoder.albert_layer_groups.0.albert_layers.0.attention", |
| "target": "encoder.albert_layer_groups.0.albert_layers.0.attention.query" |
| }, |
| { |
| "kind": "data", |
| "provenance": "intra_module", |
| "source": "encoder.albert_layer_groups.0.albert_layers.0.attention", |
| "target": "encoder.albert_layer_groups.0.albert_layers.0.attention.key" |
| }, |
| { |
| "kind": "data", |
| "provenance": "intra_module", |
| "source": "encoder.albert_layer_groups.0.albert_layers.0.attention", |
| "target": "encoder.albert_layer_groups.0.albert_layers.0.attention.value" |
| }, |
| { |
| "kind": "data", |
| "provenance": "intra_module", |
| "source": "encoder.albert_layer_groups.0.albert_layers.0.attention", |
| "target": "encoder.albert_layer_groups.0.albert_layers.0.attention.dense" |
| }, |
| { |
| "kind": "data", |
| "provenance": "observed_forward", |
| "source": "encoder", |
| "target": "pooler" |
| }, |
| { |
| "kind": "data", |
| "provenance": "observed_forward", |
| "source": "pooler", |
| "target": "pooler_activation" |
| } |
| ], |
| "extends": null, |
| "model_type": "albert", |
| "provenance": { |
| "config_class": "AlbertConfig", |
| "config_module": "transformers.models.albert.configuration_albert", |
| "model_class": "AlbertModel", |
| "model_module": "transformers.models.albert.modeling_albert" |
| }, |
| "repeats": [], |
| "schema_version": "architecture-template-v0", |
| "templates": [] |
| } |
|
|
