| { |
| "architecture": { |
| "attention_variant": "MHA", |
| "family": "causal_lm", |
| "is_moe": false, |
| "mixer": "attention", |
| "tie_word_embeddings": true, |
| "view": "decoder" |
| }, |
| "capabilities": { |
| "attention_backends": [ |
| "eager" |
| ], |
| "attention_patterns": [ |
| "causal" |
| ], |
| "attention_schedule": null, |
| "task_heads": [ |
| "causal_lm", |
| "pretraining", |
| "question_answering", |
| "sequence_classification", |
| "token_classification" |
| ], |
| "tensor_parallel": false |
| }, |
| "components": [ |
| { |
| "children": [ |
| "word_embeddings", |
| "word_embeddings_layernorm", |
| "blocks", |
| "ln_f" |
| ], |
| "class_name": "BloomModel", |
| "id": "model", |
| "kind": "model", |
| "path_pattern": "model" |
| }, |
| { |
| "attributes": { |
| "embedding_dim": "config.hidden_size", |
| "num_embeddings": "config.vocab_size", |
| "tied_lm_head": true |
| }, |
| "class_name": "Embedding", |
| "id": "word_embeddings", |
| "kind": "embedding", |
| "path_pattern": "model.word_embeddings" |
| }, |
| { |
| "attributes": { |
| "norm_type": "layer" |
| }, |
| "class_name": "LayerNorm", |
| "id": "word_embeddings_layernorm", |
| "kind": "normalization", |
| "path_pattern": "model.word_embeddings_layernorm" |
| }, |
| { |
| "attributes": { |
| "norm_type": "layer" |
| }, |
| "class_name": "LayerNorm", |
| "id": "ln_f", |
| "kind": "normalization", |
| "path_pattern": "model.ln_f" |
| } |
| ], |
| "config": { |
| "class_name": "BloomConfig", |
| "model_type": "bloom", |
| "module": "transformers.models.bloom.configuration_bloom", |
| "referenced_fields": {}, |
| "salient_fields": { |
| "hidden_size": 64, |
| "is_encoder_decoder": false, |
| "num_attention_heads": 8, |
| "num_hidden_layers": 2, |
| "tie_word_embeddings": true, |
| "vocab_size": 250880 |
| } |
| }, |
| "dataflow": { |
| "input": { |
| "name": "input_ids", |
| "shape": [ |
| "B", |
| "S" |
| ] |
| }, |
| "output": { |
| "shape": [ |
| "B", |
| "S", |
| "config.hidden_size" |
| ] |
| }, |
| "shapes": { |
| "block.input_layernorm": { |
| "in": [ |
| "B", |
| "S", |
| "config.hidden_size" |
| ], |
| "out": [ |
| "B", |
| "S", |
| "config.hidden_size" |
| ] |
| }, |
| "block.mlp": { |
| "in": [ |
| "B", |
| "S", |
| "config.hidden_size" |
| ], |
| "out": [ |
| "B", |
| "S", |
| "config.hidden_size" |
| ] |
| }, |
| "block.post_attention_layernorm": { |
| "in": [ |
| "B", |
| "S", |
| "config.hidden_size" |
| ], |
| "out": [ |
| "B", |
| "S", |
| "config.hidden_size" |
| ] |
| }, |
| "block.self_attention": { |
| "in": [ |
| "B", |
| "S", |
| "config.hidden_size" |
| ], |
| "out": [ |
| "B", |
| "S", |
| "config.hidden_size" |
| ] |
| }, |
| "blocks": { |
| "in": [ |
| "B", |
| "S", |
| "config.hidden_size" |
| ], |
| "out": [ |
| "B", |
| "S", |
| "config.hidden_size" |
| ] |
| }, |
| "ln_f": { |
| "in": [ |
| "B", |
| "S", |
| "config.hidden_size" |
| ], |
| "out": [ |
| "B", |
| "S", |
| "config.hidden_size" |
| ] |
| }, |
| "word_embeddings": { |
| "in": [ |
| "B", |
| "S" |
| ], |
| "out": [ |
| "B", |
| "S", |
| "config.hidden_size" |
| ] |
| }, |
| "word_embeddings_layernorm": { |
| "in": [ |
| "B", |
| "S", |
| "config.hidden_size" |
| ], |
| "out": [ |
| "B", |
| "S", |
| "config.hidden_size" |
| ] |
| } |
| }, |
| "source": "observed_forward_meta" |
| }, |
| "edges": [ |
| { |
| "kind": "data", |
| "source": "word_embeddings", |
| "target": "word_embeddings_layernorm" |
| }, |
| { |
| "kind": "data", |
| "source": "word_embeddings_layernorm", |
| "target": "blocks" |
| }, |
| { |
| "kind": "data", |
| "source": "blocks", |
| "target": "ln_f" |
| }, |
| { |
| "kind": "mask", |
| "source": "input:attention_mask", |
| "target": "block.self_attention" |
| }, |
| { |
| "kind": "residual", |
| "source": "block", |
| "target": "block.self_attention" |
| }, |
| { |
| "kind": "residual", |
| "source": "block", |
| "target": "block.mlp" |
| }, |
| { |
| "kind": "data", |
| "provenance": "intra_module", |
| "source": "block.self_attention", |
| "target": "block.self_attention.query_key_value" |
| }, |
| { |
| "kind": "data", |
| "provenance": "intra_module", |
| "source": "block.self_attention", |
| "target": "block.self_attention.dense" |
| }, |
| { |
| "kind": "data", |
| "provenance": "intra_module", |
| "source": "block.mlp", |
| "target": "block.mlp.dense_h_to_4h" |
| }, |
| { |
| "kind": "data", |
| "provenance": "intra_module", |
| "source": "block.mlp.dense_h_to_4h", |
| "target": "block.mlp.dense_4h_to_h" |
| }, |
| { |
| "kind": "cache_read", |
| "source": "state:kv_cache", |
| "target": "block.self_attention" |
| }, |
| { |
| "kind": "cache_write", |
| "source": "block.self_attention", |
| "target": "state:kv_cache" |
| }, |
| { |
| "kind": "data", |
| "provenance": "observed_forward", |
| "source": "block.input_layernorm", |
| "target": "block.self_attention" |
| }, |
| { |
| "kind": "data", |
| "provenance": "observed_forward", |
| "source": "block.self_attention", |
| "target": "block.post_attention_layernorm" |
| }, |
| { |
| "kind": "data", |
| "provenance": "observed_forward", |
| "source": "block.post_attention_layernorm", |
| "target": "block.mlp" |
| } |
| ], |
| "extends": null, |
| "model_type": "bloom", |
| "provenance": { |
| "config_class": "BloomConfig", |
| "config_module": "transformers.models.bloom.configuration_bloom", |
| "model_class": "BloomModel", |
| "model_module": "transformers.models.bloom.modeling_bloom" |
| }, |
| "repeats": [ |
| { |
| "body": "block", |
| "container_path_pattern": "model.h", |
| "count": 2, |
| "count_expr": "config.num_hidden_layers", |
| "count_source": "config", |
| "id": "blocks", |
| "index_symbol": "i", |
| "item_path_pattern": "model.h.{i}", |
| "kind": "symbolic_repeat", |
| "provenance": { |
| "class_name": "BloomBlock", |
| "container_path_pattern": "model.h", |
| "item_path_pattern": "model.h.{i}", |
| "source": "module_tree_repeat_collapse" |
| }, |
| "repeated_class_name": "BloomBlock" |
| } |
| ], |
| "schema_version": "architecture-template-v0", |
| "templates": [ |
| { |
| "children": [ |
| "block.input_layernorm", |
| "block.self_attention", |
| "block.post_attention_layernorm", |
| "block.mlp" |
| ], |
| "class_name": "BloomBlock", |
| "id": "block", |
| "kind": "transformer_block", |
| "path_pattern": "model.h.{i}" |
| }, |
| { |
| "attributes": { |
| "norm_type": "layer" |
| }, |
| "class_name": "LayerNorm", |
| "id": "block.input_layernorm", |
| "kind": "normalization", |
| "path_pattern": "model.h.{i}.input_layernorm" |
| }, |
| { |
| "attributes": { |
| "head_dim": 8, |
| "n_heads": 8, |
| "n_kv_heads": 8, |
| "pattern": "causal", |
| "rope": false, |
| "variant": "MHA" |
| }, |
| "children": [ |
| "block.self_attention.query_key_value", |
| "block.self_attention.dense" |
| ], |
| "class_name": "BloomAttention", |
| "id": "block.self_attention", |
| "kind": "attention", |
| "path_pattern": "model.h.{i}.self_attention" |
| }, |
| { |
| "attributes": { |
| "in_features": "config.hidden_size", |
| "out_features": 192 |
| }, |
| "class_name": "Linear", |
| "id": "block.self_attention.query_key_value", |
| "kind": "projection", |
| "path_pattern": "model.h.{i}.self_attention.query_key_value" |
| }, |
| { |
| "attributes": { |
| "in_features": "config.hidden_size", |
| "out_features": "config.hidden_size" |
| }, |
| "class_name": "Linear", |
| "id": "block.self_attention.dense", |
| "kind": "projection", |
| "path_pattern": "model.h.{i}.self_attention.dense" |
| }, |
| { |
| "attributes": { |
| "norm_type": "layer" |
| }, |
| "class_name": "LayerNorm", |
| "id": "block.post_attention_layernorm", |
| "kind": "normalization", |
| "path_pattern": "model.h.{i}.post_attention_layernorm" |
| }, |
| { |
| "attributes": { |
| "hidden_size": 64 |
| }, |
| "children": [ |
| "block.mlp.dense_h_to_4h", |
| "block.mlp.gelu_impl", |
| "block.mlp.dense_4h_to_h" |
| ], |
| "class_name": "BloomMLP", |
| "id": "block.mlp", |
| "kind": "feed_forward", |
| "path_pattern": "model.h.{i}.mlp" |
| }, |
| { |
| "attributes": { |
| "in_features": "config.hidden_size", |
| "out_features": 256 |
| }, |
| "class_name": "Linear", |
| "id": "block.mlp.dense_h_to_4h", |
| "kind": "projection", |
| "path_pattern": "model.h.{i}.mlp.dense_h_to_4h" |
| }, |
| { |
| "attributes": { |
| "hidden_size": 64 |
| }, |
| "class_name": "BloomGelu", |
| "id": "block.mlp.gelu_impl", |
| "kind": "feed_forward", |
| "path_pattern": "model.h.{i}.mlp.gelu_impl" |
| }, |
| { |
| "attributes": { |
| "in_features": 256, |
| "out_features": "config.hidden_size" |
| }, |
| "class_name": "Linear", |
| "id": "block.mlp.dense_4h_to_h", |
| "kind": "projection", |
| "path_pattern": "model.h.{i}.mlp.dense_4h_to_h" |
| } |
| ] |
| } |
|
|
