Custom Dataset

[zzzmm1]

Thanks for your excellent work!
I would like to know how to pretrain and finetune GraphGPT in our custom dataset, and how to determine the tokenizer?
Any suggestion would be helpful.

[zhaoqf123]

There are 3 steps to customize for your own datasets.

1. Create and register your custom data in script src/data/data_sources.py.

The details are as follows:

If your graph dataset contains many small graphs, and the task is graph-level, you can refer to ogbg-molpcba/reddit_threads.

If your graph dataset contains only one big graph, and the task is edge-level or node-level, you can refer to ogbl-ppa/ogbl-citation2 for edge-level, and refer to ogbn-proteins/ogbn-arxiv for node-level.

By the way, if your raw dataset is not of the format InMemoryDataset from torch_geometric package, you will need to convert to this format first.

2. Create and modify the tokenization config file.

Base on your task type, you can refer to tokenization config in the corresponding folders in examples.

Usually, the entries you need to modify in the json file includes dataset(the name of the function you register in src/data/data_sources.py), attr_world_identifier, semantics->node->dim, and semantics->edge->dim.

3. Create and modify the pre-train/fine-tune bash scripts

Base on your task type, you can refer to bash scripts xxx_pretrain.sh and xxx_supervised.sh in the corresponding folders in examples.

The parameters you need to modify usually includes dataset_name, dataset_source (usually the same), tokenization_config_file, ds_prefix and etc.

The vocabulary will be created automatically for this specific dataset.

Before implementing your own dataset, please run the toy_examples and make sure no bugs.

Finally, if your graph dataset is very different from the dataset from datasets in the examples folder, you may send some samples to us and we will try to help figure out how to implement it.

[xylyu-zju]

I am also using a custom dataset, refer to the format of ogbg-molpcba, my dataset has 3350 items, it is single-label classification, when I run the new dataset I always get the same error, the following is displayed, I want to know if it is related to the virtual node? If so, how should I disable it? If not, how should I solve it?
’‘’
[2025-06-28 17:07:18.221951] NOT RESET samples of GraphsMapDataset of 335 graphs for epoch None!

idx_tuple: None
[Warning] permute_nodes enabled! edge_attr remains the same; edge_index and node-attrs will be affected!
[Warning] permute_nodes enabled! edge_attr remains the same; edge_index and node-attrs will be affected!
[Warning] permute_nodes enabled! edge_attr remains the same; edge_index and node-attrs will be affected!
[Warning] permute_nodes enabled! edge_attr remains the same; edge_index and node-attrs will be affected!
[Warning] permute_nodes enabled! edge_attr remains the same; edge_index and node-attrs will be affected!

[2025-06-28 17:07:18.222320] NOT RESET samples of GraphsMapDataset of 335 graphs for epoch None!
idx_tuple: None
Split dataset based on given train/valid/test index!
Train: 2680, Valid: 335, Test: 335!
2680
(0, Data(edge_index=[2, 32], edge_attr=[32, 1], x=[8, 6], y=[1, 1], num_nodes=8, idx=0, idx_of_ds=0))
335
Traceback (most recent call last):
File "./examples/train_supervised.py", line 845, in
fire.Fire(train)
File "/home/xylv@iaclab.ic/anaconda3/envs/graphgpt/lib/python3.8/site-packages/fire/core.py", line 141, in Fire
component_trace = _Fire(component, args, parsed_flag_args, context, name)
File "/home/xylv@iaclab.ic/anaconda3/envs/graphgpt/lib/python3.8/site-packages/fire/core.py", line 475, in _Fire
component, remaining_args = _CallAndUpdateTrace(
File "/home/xylv@iaclab.ic/anaconda3/envs/graphgpt/lib/python3.8/site-packages/fire/core.py", line 691, in _CallAndUpdateTrace
component = fn(*varargs, **kwargs)
File "./examples/train_supervised.py", line 314, in train
print(dataset[idx])
File "/home/xylv@iaclab.ic/python_code/analog-gpt/./src/data/dataset_map.py", line 2818, in getitem
graph, _ = nx_utils.permute_nodes(graph, self.g)
File "/home/xylv@iaclab.ic/python_code/analog-gpt/./src/utils/nx_utils.py", line 1902, in permute_nodes
new_graph.edge_index = permu[graph.edge_index]
IndexError: index 54 is out of bounds for dimension 0 with size 20
‘’‘

[xylyu-zju]

When I close it"permute_nodes: bool = False,",new error comes:

permute_nodes: bool = False,

Inspecting tokenization results!
Tokenize graph:
Data(edge_index=[2, 32], edge_attr=[32, 1], x=[8, 6], y=[1, 1], num_nodes=8, idx=0, idx_of_ds=0)
Tokens:
[['190', 'molecule#node#0#218', 'molecule#node#1#0', 'molecule#node#2#0', 'molecule#node#3#0', 'molecule#node#4#1', 'molecule#node#5#0'],
 ['191', 'molecule#node#0#221', 'molecule#node#1#1', 'molecule#node#2#1', 'molecule#node#3#0', 'molecule#node#4#0', 'molecule#node#5#0'],
 ['192', 'molecule#node#0#0', 'molecule#node#1#0', 'molecule#node#2#1', 'molecule#node#3#0', 'molecule#node#4#0', 'molecule#node#5#0'],
 ['193', 'molecule#node#0#221', 'molecule#node#1#1', 'molecule#node#2#1', 'molecule#node#3#0', 'molecule#node#4#0', 'molecule#node#5#0'],
 ['190', 'molecule#node#0#218', 'molecule#node#1#0', 'molecule#node#2#0', 'molecule#node#3#0', 'molecule#node#4#1', 'molecule#node#5#0'],
 ['<eos>', '<eos>', '<eos>', '<eos>', '<eos>', '<eos>', '<eos>']]
Labels:
[['191', 'molecule#node#0#221', 'molecule#node#1#1', 'molecule#node#2#1', 'molecule#node#3#0', 'molecule#node#4#0', 'molecule#node#5#0'],
 ['192', 'molecule#node#0#0', 'molecule#node#1#0', 'molecule#node#2#1', 'molecule#node#3#0', 'molecule#node#4#0', 'molecule#node#5#0'],
 ['193', 'molecule#node#0#221', 'molecule#node#1#1', 'molecule#node#2#1', 'molecule#node#3#0', 'molecule#node#4#0', 'molecule#node#5#0'],
 ['190', 'molecule#node#0#218', 'molecule#node#1#0', 'molecule#node#2#0', 'molecule#node#3#0', 'molecule#node#4#1', 'molecule#node#5#0'],
 ['<eos>', '<eos>', '<eos>', '<eos>', '<eos>', '<eos>', '<eos>'],
 ['<eos>', '<eos>', '<eos>', '<eos>', '<eos>', '<eos>', '<eos>']]
embed:[]
Traceback (most recent call last):
  File "./examples/train_supervised.py", line 845, in <module>
    fire.Fire(train)
  File "/home/xylv@iaclab.ic/anaconda3/envs/graphgpt/lib/python3.8/site-packages/fire/core.py", line 141, in Fire
    component_trace = _Fire(component, args, parsed_flag_args, context, name)
  File "/home/xylv@iaclab.ic/anaconda3/envs/graphgpt/lib/python3.8/site-packages/fire/core.py", line 475, in _Fire
    component, remaining_args = _CallAndUpdateTrace(
  File "/home/xylv@iaclab.ic/anaconda3/envs/graphgpt/lib/python3.8/site-packages/fire/core.py", line 691, in _CallAndUpdateTrace
    component = fn(*varargs, **kwargs)
  File "./examples/train_supervised.py", line 324, in train
    inspect_tokenization_results(train_dataset, gtokenizer)
  File "/home/xylv@iaclab.ic/python_code/analog-gpt/./src/utils/inspection_utils.py", line 186, in inspect_tokenization_results
    in_dict = gtokenizer.convert_tokens_to_ids(tokens, labels)
  File "/home/xylv@iaclab.ic/python_code/analog-gpt/./src/data/tokenizer.py", line 3639, in convert_tokens_to_ids
    seq_tokens_id = [
  File "/home/xylv@iaclab.ic/python_code/analog-gpt/./src/data/tokenizer.py", line 3640, in <listcomp>
    [self.vocab_map[token] for token in feat_tokens]
  File "/home/xylv@iaclab.ic/python_code/analog-gpt/./src/data/tokenizer.py", line 3640, in <listcomp>
    [self.vocab_map[token] for token in feat_tokens]
KeyError: 'molecule#node#0#218'

My config setting is:

  "name_or_path": "ogbg_molpcba",
  "dataset": "ogbg-molpcba",
  "attr_world_identifier": "molecule",
  "task_type": "pretrain-mlm",
  "pretrain_mlm": {
    "name": "polynomial",
    "params": {"fixed_ratio": 0.7, "power": 1, "mtp": [0.8, 0, 0.2]},
    "info": "name->polynomial|cosine|fixed,power->3/2/1/0.5"
  },
  "vocab_file": "vocab",
  "sampling": null,
  "semantics": {
    "attr_assignment": "first",
    "attr_shuffle": false,
    "node": {
      "discrete": "x",
      "continuous": null,
      "ignored_val": "0",
      "dim": 6
    },
    "edge": {
      "discrete":null,
      "continuous": null,
      "ignored_val": null,
      "dim": 0
    },
    "graph": {
      "discrete": null,
      "continuous": null,
      "ignored_val": null
    },
    "common": {
      "reserved_token": [
        "semantics_0",
        "semantics_1",
        "semantics_2",
        "semantics_3",
        "semantics_4",
        "semantics_5",
        "semantics_6",
        "semantics_7",
        "semantics_8",
        "semantics_9"
      ],
      "numbers": [
        "<e>",
        "<.>",
        "<->",
        "<0>",
        "<1>",
        "<2>",
        "<3>",
        "<4>",
        "<5>",
        "<6>",
        "<7>",
        "<8>",
        "<9>"
      ]
    },
    "instructions": {
      "enable": false
    }
  },
  "structure": {
    "nx": {
      "enable": false
    },
    "node": {
      "bos_token": "<bos>",
      "eos_token": "<eos>",
      "new_node_token": "<new>",
      "node_scope": 512,
      "scope_base": 512,
      "cyclic": 1
    },
    "edge": {
      "remove_edge_type_token": true,

      "in_token": "<edge_in>",
      "out_token":"<edge_out>",
      "bi_token": "<edge_bi>",
      "jump_token": "<edge_jump>"
    },
    "graph": {
      "summary_token": "<gsum>"
    },
    "common": {
      "mask_token": "<mask>",
      "icl_token": "<icl>",
      "sep_token": "<sep>",
      "reserved_token": [
        "structure_0",
        "structure_1",
        "structure_2",
        "structure_3",
        "structure_4",
        "structure_5",
        "structure_6",
        "structure_7",
        "structure_8",
        "structure_9"
      ]
    }
  },
  "tokenizer_class": "StackedGSTTokenizer"
}

And the vocab is automatically generated by vocab_builder.py. I want to know how to solve this, thanks!

[zhaoqf123]

It looks like your dataset is not molecule, or not in the conventional format as in ogbl-molpcba if it is actually molecules.

Please delete the vocab and use the config below, and the vocab file shall be created automatically.
The main change of json elements are: attr_world_identifier, semantics->edge->discrete|dim:

  "name_or_path": "ogbg_molpcba",
  "dataset": "ogbg-molpcba",
  "attr_world_identifier": "your-dataset-name",
  "task_type": "pretrain-mlm",
  "pretrain_mlm": {
    "name": "polynomial",
    "params": {"fixed_ratio": 0.7, "power": 1, "mtp": [0.8, 0, 0.2]},
    "info": "name->polynomial|cosine|fixed,power->3/2/1/0.5"
  },
  "vocab_file": "vocab",
  "sampling": null,
  "semantics": {
    "attr_assignment": "first",
    "attr_shuffle": false,
    "node": {
      "discrete": "x",
      "continuous": null,
      "ignored_val": "0",
      "dim": 6
    },
    "edge": {
      "discrete":"edge_attr",
      "continuous": null,
      "ignored_val": null,
      "dim": 1
    },
    "graph": {
      "discrete": null,
      "continuous": null,
      "ignored_val": null
    },
    "common": {
      "reserved_token": [
        "semantics_0",
        "semantics_1",
        "semantics_2",
        "semantics_3",
        "semantics_4",
        "semantics_5",
        "semantics_6",
        "semantics_7",
        "semantics_8",
        "semantics_9"
      ],
      "numbers": [
        "<e>",
        "<.>",
        "<->",
        "<0>",
        "<1>",
        "<2>",
        "<3>",
        "<4>",
        "<5>",
        "<6>",
        "<7>",
        "<8>",
        "<9>"
      ]
    },
    "instructions": {
      "enable": false
    }
  },
  "structure": {
    "nx": {
      "enable": false
    },
    "node": {
      "bos_token": "<bos>",
      "eos_token": "<eos>",
      "new_node_token": "<new>",
      "node_scope": 512,
      "scope_base": 512,
      "cyclic": 1
    },
    "edge": {
      "remove_edge_type_token": true,

      "in_token": "<edge_in>",
      "out_token":"<edge_out>",
      "bi_token": "<edge_bi>",
      "jump_token": "<edge_jump>"
    },
    "graph": {
      "summary_token": "<gsum>"
    },
    "common": {
      "mask_token": "<mask>",
      "icl_token": "<icl>",
      "sep_token": "<sep>",
      "reserved_token": [
        "structure_0",
        "structure_1",
        "structure_2",
        "structure_3",
        "structure_4",
        "structure_5",
        "structure_6",
        "structure_7",
        "structure_8",
        "structure_9"
      ]
    }
  },
  "tokenizer_class": "StackedGSTTokenizer"
}

Besides, from your log, there seems to be some problem with the data:

Inspecting tokenization results!
Tokenize graph:
Data(edge_index=[2, 32], edge_attr=[32, 1], x=[8, 6], y=[1, 1], num_nodes=8, idx=0, idx_of_ds=0)
Tokens:
[['190', 'molecule#node#0#218', 'molecule#node#1#0', 'molecule#node#2#0', 'molecule#node#3#0', 'molecule#node#4#1', 'molecule#node#5#0'],
 ['191', 'molecule#node#0#221', 'molecule#node#1#1', 'molecule#node#2#1', 'molecule#node#3#0', 'molecule#node#4#0', 'molecule#node#5#0'],
 ['192', 'molecule#node#0#0', 'molecule#node#1#0', 'molecule#node#2#1', 'molecule#node#3#0', 'molecule#node#4#0', 'molecule#node#5#0'],
 ['193', 'molecule#node#0#221', 'molecule#node#1#1', 'molecule#node#2#1', 'molecule#node#3#0', 'molecule#node#4#0', 'molecule#node#5#0'],
 ['190', 'molecule#node#0#218', 'molecule#node#1#0', 'molecule#node#2#0', 'molecule#node#3#0', 'molecule#node#4#1', 'molecule#node#5#0'],
 ['<eos>', '<eos>', '<eos>', '<eos>', '<eos>', '<eos>', '<eos>']]

The graph has 8 nodes, but resulting Eulerian sequence 190-191-192-193-190 implies only 4 nodes are presented. It's wired.

Please use the config to run again. If got problem, please post here.

Please also take a look at the log above even if the code can run through. :)

[xylyu-zju]

Thank you for your guidance. After I modified it, I still get the idx error. Can I send you a dataset and you give me some suggestions? My dataset is a circuit netlist dataset, with 15 labels and 3350 netlist-generated graphs. There are node features, and the edges do not contain information except for the connection relationship. Each graph corresponds to a label. Before this, I have successfully run toy-example and ogbg-molpcba.

Inspecting graph of index 0
Inspecting tokenization results!
Tokenize graph:
Data(edge_index=[2, 32], edge_attr=[32, 1], x=[8, 6], y=[1, 1], num_nodes=8, idx=0, idx_of_ds=0)

Tokens:
[['20', 'ogbg_analog#node#0#219', 'ogbg_analog#node#1#0', 'ogbg_analog#node#2#0', 'ogbg_analog#node#3#0', 'ogbg_analog#node#4#1', 'ogbg_analog#node#5#0'],
 ['21', 'ogbg_analog#node#0#221', 'ogbg_analog#node#1#1', 'ogbg_analog#node#2#1', 'ogbg_analog#node#3#0', 'ogbg_analog#node#4#0', 'ogbg_analog#node#5#0'],
 ['22', 'ogbg_analog#node#0#0', 'ogbg_analog#node#1#0', 'ogbg_analog#node#2#1', 'ogbg_analog#node#3#0', 'ogbg_analog#node#4#0', 'ogbg_analog#node#5#0'],
 ['23', 'ogbg_analog#node#0#221', 'ogbg_analog#node#1#1', 'ogbg_analog#node#2#1', 'ogbg_analog#node#3#0', 'ogbg_analog#node#4#0', 'ogbg_analog#node#5#0'],
 ['24', 'ogbg_analog#node#0#221', 'ogbg_analog#node#1#0', 'ogbg_analog#node#2#0', 'ogbg_analog#node#3#0', 'ogbg_analog#node#4#1', 'ogbg_analog#node#5#0'],
 ['23', 'ogbg_analog#node#0#221', 'ogbg_analog#node#1#1', 'ogbg_analog#node#2#1', 'ogbg_analog#node#3#0', 'ogbg_analog#node#4#0', 'ogbg_analog#node#5#0'],
 ['22', 'ogbg_analog#node#0#0', 'ogbg_analog#node#1#0', 'ogbg_analog#node#2#1', 'ogbg_analog#node#3#0', 'ogbg_analog#node#4#0', 'ogbg_analog#node#5#0'],
 ['25', 'ogbg_analog#node#0#221', 'ogbg_analog#node#1#1', 'ogbg_analog#node#2#1', 'ogbg_analog#node#3#0', 'ogbg_analog#node#4#0', 'ogbg_analog#node#5#0'],
 ['26', 'ogbg_analog#node#0#218', 'ogbg_analog#node#1#0', 'ogbg_analog#node#2#0', 'ogbg_analog#node#3#0', 'ogbg_analog#node#4#1', 'ogbg_analog#node#5#0'],
 ['27', 'ogbg_analog#node#0#221', 'ogbg_analog#node#1#1', 'ogbg_analog#node#2#1', 'ogbg_analog#node#3#0', 'ogbg_analog#node#4#0', 'ogbg_analog#node#5#0'],
 ['22', 'ogbg_analog#node#0#0', 'ogbg_analog#node#1#0', 'ogbg_analog#node#2#1', 'ogbg_analog#node#3#0', 'ogbg_analog#node#4#0', 'ogbg_analog#node#5#0'],
 ['<eos>', '<eos>', '<eos>', '<eos>', '<eos>', '<eos>', '<eos>']]
Labels:
[['21', 'ogbg_analog#node#0#221', 'ogbg_analog#node#1#1', 'ogbg_analog#node#2#1', 'ogbg_analog#node#3#0', 'ogbg_analog#node#4#0', 'ogbg_analog#node#5#0'],
 ['22', 'ogbg_analog#node#0#0', 'ogbg_analog#node#1#0', 'ogbg_analog#node#2#1', 'ogbg_analog#node#3#0', 'ogbg_analog#node#4#0', 'ogbg_analog#node#5#0'],
 ['23', 'ogbg_analog#node#0#221', 'ogbg_analog#node#1#1', 'ogbg_analog#node#2#1', 'ogbg_analog#node#3#0', 'ogbg_analog#node#4#0', 'ogbg_analog#node#5#0'],
 ['24', 'ogbg_analog#node#0#221', 'ogbg_analog#node#1#0', 'ogbg_analog#node#2#0', 'ogbg_analog#node#3#0', 'ogbg_analog#node#4#1', 'ogbg_analog#node#5#0'],
 ['23', 'ogbg_analog#node#0#221', 'ogbg_analog#node#1#1', 'ogbg_analog#node#2#1', 'ogbg_analog#node#3#0', 'ogbg_analog#node#4#0', 'ogbg_analog#node#5#0'],
 ['22', 'ogbg_analog#node#0#0', 'ogbg_analog#node#1#0', 'ogbg_analog#node#2#1', 'ogbg_analog#node#3#0', 'ogbg_analog#node#4#0', 'ogbg_analog#node#5#0'],
 ['25', 'ogbg_analog#node#0#221', 'ogbg_analog#node#1#1', 'ogbg_analog#node#2#1', 'ogbg_analog#node#3#0', 'ogbg_analog#node#4#0', 'ogbg_analog#node#5#0'],
 ['26', 'ogbg_analog#node#0#218', 'ogbg_analog#node#1#0', 'ogbg_analog#node#2#0', 'ogbg_analog#node#3#0', 'ogbg_analog#node#4#1', 'ogbg_analog#node#5#0'],
 ['27', 'ogbg_analog#node#0#221', 'ogbg_analog#node#1#1', 'ogbg_analog#node#2#1', 'ogbg_analog#node#3#0', 'ogbg_analog#node#4#0', 'ogbg_analog#node#5#0'],
 ['22', 'ogbg_analog#node#0#0', 'ogbg_analog#node#1#0', 'ogbg_analog#node#2#1', 'ogbg_analog#node#3#0', 'ogbg_analog#node#4#0', 'ogbg_analog#node#5#0'],
 ['<eos>', '<eos>', '<eos>', '<eos>', '<eos>', '<eos>', '<eos>'],
 ['<eos>', '<eos>', '<eos>', '<eos>', '<eos>', '<eos>', '<eos>']]
embed:[]

Tokenized results:
{'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
 'embed': array([], shape=(12, 0), dtype=float64),
 'input_ids': [[42, 659, 735, 737, 739, 742, 743],
               [43, 661, 736, 738, 739, 741, 743],
               [44, 563, 735, 738, 739, 741, 743],
               [45, 661, 736, 738, 739, 741, 743],
               [46, 661, 735, 737, 739, 742, 743],
               [45, 661, 736, 738, 739, 741, 743],
               [44, 563, 735, 738, 739, 741, 743],
               [47, 661, 736, 738, 739, 741, 743],
               [48, 658, 735, 737, 739, 742, 743],
               [49, 661, 736, 738, 739, 741, 743],
               [44, 563, 735, 738, 739, 741, 743],
               [19, 19, 19, 19, 19, 19, 19]],
 'labels': [[43, 661, 736, 738, 739, 741, 743],
            [44, 563, 735, 738, 739, 741, 743],
            [45, 661, 736, 738, 739, 741, 743],
            [46, 661, 735, 737, 739, 742, 743],
            [45, 661, 736, 738, 739, 741, 743],
            [44, 563, 735, 738, 739, 741, 743],
            [47, 661, 736, 738, 739, 741, 743],
            [48, 658, 735, 737, 739, 742, 743],
            [49, 661, 736, 738, 739, 741, 743],
            [44, 563, 735, 738, 739, 741, 743],
            [19, 19, 19, 19, 19, 19, 19],
            [19, 19, 19, 19, 19, 19, 19]],
 'position_ids': [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]}

Inputs for model:
{'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
 'embed': array([], shape=(12, 0), dtype=float64),
 'graph_labels': 1,
 'input_ids': [[42, 659, 735, 737, 739, 742, 743],
               [43, 661, 736, 738, 739, 741, 743],
               [44, 563, 735, 738, 739, 741, 743],
               [45, 661, 736, 738, 739, 741, 743],
               [46, 661, 735, 737, 739, 742, 743],
               [45, 661, 736, 738, 739, 741, 743],
               [44, 563, 735, 738, 739, 741, 743],
               [47, 661, 736, 738, 739, 741, 743],
               [48, 658, 735, 737, 739, 742, 743],
               [49, 661, 736, 738, 739, 741, 743],
               [44, 563, 735, 738, 739, 741, 743],
               [19, 19, 19, 19, 19, 19, 19]],
 'labels': [[43, 661, 736, 738, 739, 741, 743],
            [44, 563, 735, 738, 739, 741, 743],
            [45, 661, 736, 738, 739, 741, 743],
            [46, 661, 735, 737, 739, 742, 743],
            [45, 661, 736, 738, 739, 741, 743],
            [44, 563, 735, 738, 739, 741, 743],
            [47, 661, 736, 738, 739, 741, 743],
            [48, 658, 735, 737, 739, 742, 743],
            [49, 661, 736, 738, 739, 741, 743],
            [44, 563, 735, 738, 739, 741, 743],
            [19, 19, 19, 19, 19, 19, 19],
            [19, 19, 19, 19, 19, 19, 19]],
 'position_ids': [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]}

[Warning] Set eos_idx to 100000000 for task graph!
steps_per_epoch: 20

total_num_steps: 640
warmup_num_steps: 192


Final model config for supervised task:
GraphGPTConfig {
  "attention_bias": false,
  "attention_dropout": 0.1,
  "bos_token_id": 20,
  "causal_attention": false,
  "cls_token_id": null,
  "dropout": 0,
  "embed_dim": 0,
  "embed_pdrop": 0,
  "eos_token_id": 19,
  "hidden_act": "gelu",
  "hidden_size": 128,
  "id2label": {
    "0": "LABEL_0",
    "1": "LABEL_1",
    "2": "LABEL_2",
    "3": "LABEL_3",
    "4": "LABEL_4",
    "5": "LABEL_5",
    "6": "LABEL_6",
    "7": "LABEL_7",
    "8": "LABEL_8",
    "9": "LABEL_9",
    "10": "LABEL_10",
    "11": "LABEL_11",
    "12": "LABEL_12",
    "13": "LABEL_13",
    "14": "LABEL_14"
  },
  "initializer_range": 0.02,
  "intermediate_size": 512,
  "label2id": {
    "LABEL_0": 0,
    "LABEL_1": 1,
    "LABEL_10": 10,
    "LABEL_11": 11,
    "LABEL_12": 12,
    "LABEL_13": 13,
    "LABEL_14": 14,
    "LABEL_2": 2,
    "LABEL_3": 3,
    "LABEL_4": 4,
    "LABEL_5": 5,
    "LABEL_6": 6,
    "LABEL_7": 7,
    "LABEL_8": 8,
    "LABEL_9": 9
  },
  "layer_scale_init_value": 1,
  "loss_type": null,
  "max_position_embeddings": 1024,
  "mlp": [],
  "mlp_pdrop": 0,
  "model_type": "graphgpt",
  "next_n_token": 1,
  "num_attention_heads": 2,
  "num_hidden_layers": 2,
  "num_key_value_heads": 2,
  "num_neg": null,
  "pad_token_id": 0,
  "path_pdrop": 0.1,
  "pooling_method": "last",
  "pretraining_tp": 1,
  "problem_type": "single_label_classification",
  "rms_norm_eps": 1e-06,
  "rope_scaling": null,
  "rope_theta": 10000,
  "stack_method": "short",
  "stacked_feat": 7,
  "stacked_feat_agg_method": "gated",
  "tie_word_embeddings": false,
  "transformers_version": "4.38.2",
  "use_cache": true,
  "use_ntp": false,
  "vocab_size": 745
}


[2025-06-30 01:11:14,809] [INFO] [comm.py:652:init_distributed] cdb=None
Applying dropout in backbone transformer

Monkey Patch LlamaModel's method `_update_causal_mask`!

Pooling in last layer is last!
trainable params: 623616 || all params: 623616 || trainable%: 100.0
GraphGPTDoubleHeadsModel(
  (model): LlamaModel(
    (embed_tokens): Embedding(745, 128, padding_idx=0)
    (norm): LlamaRMSNorm()
    (layers): ModuleList(
      (0): LlamaDecoderLayer(
        (self_attn): LlamaSdpaAttention(
          (q_proj): Linear(in_features=128, out_features=128, bias=False)
          (k_proj): Linear(in_features=128, out_features=128, bias=False)
          (v_proj): Linear(in_features=128, out_features=128, bias=False)
          (o_proj): Linear(in_features=128, out_features=128, bias=False)
          (rotary_emb): LlamaRotaryEmbedding()
        )
        (mlp): LlamaMLP(
          (gate_proj): Linear(in_features=128, out_features=512, bias=False)
          (up_proj): Linear(in_features=128, out_features=512, bias=False)
          (down_proj): Linear(in_features=512, out_features=128, bias=False)
          (act_fn): GELUActivation()
        )
        (input_layernorm): LlamaRMSNorm()
        (post_attention_layernorm): LlamaRMSNorm()
        (drop_path): Identity()
      )
      (1): LlamaDecoderLayer(
        (self_attn): LlamaSdpaAttention(
          (q_proj): Linear(in_features=128, out_features=128, bias=False)
          (k_proj): Linear(in_features=128, out_features=128, bias=False)
          (v_proj): Linear(in_features=128, out_features=128, bias=False)
          (o_proj): Linear(in_features=128, out_features=128, bias=False)
          (rotary_emb): LlamaRotaryEmbedding()
        )
        (mlp): LlamaMLP(
          (gate_proj): Linear(in_features=128, out_features=512, bias=False)
          (up_proj): Linear(in_features=128, out_features=512, bias=False)
          (down_proj): Linear(in_features=512, out_features=128, bias=False)
          (act_fn): GELUActivation()
        )
        (input_layernorm): LlamaRMSNorm()
        (post_attention_layernorm): LlamaRMSNorm()
        (drop_path): DropPath(p=0.1)
      )
    )
  )
  (stacked_feat_agg): StackedFeatAggregation(stacked_feat=7, hidden_size=128)
  (score): Linear(in_features=128, out_features=15, bias=False)
)

ds_config:
{'activation_checkpointing': {'contiguous_memory_optimization': True,
                              'cpu_checkpointing': False,
                              'partition_activations': True},
 'flops_profiler': {'detailed': True,
                    'enabled': False,
                    'module_depth': -1,
                    'output_file': None,
                    'profile_step': 1,
                    'top_modules': 1},
 'fp16': {'enabled': False,
          'hysteresis': 2,
          'initial_scale_power': 16,
          'loss_scale': 0,
          'loss_scale_window': 1000,
          'min_loss_scale': 1},
 'gradient_accumulation_steps': 1,
 'gradient_clipping': 1,
 'optimizer': {'params': {'betas': '[0.9,0.99]',
                          'bias_correction': True,
                          'eps': 1e-10,
                          'lr': 0.001,
                          'weight_decay': 0},
               'type': 'Adam'},
 'scheduler': {'params': {}, 'type': 'OneCycleLR'},
 'train_batch_size': 128,
 'train_micro_batch_size_per_gpu': 128,
 'wall_clock_breakdown': False,
 'zero_optimization': {'overlap_comm': True, 'stage': 2}}
Pop scheduler in ds_config: {'type': 'OneCycleLR', 'params': {'max_lr': 0.001, 'min_lr': 4e-09, 'div_factor': 25, 'final_div_factor': 10000.0, 'total_steps': 640, 'pct_start': 0.3, 'last_epoch': -1}}
[2025-06-30 01:11:14,851] [INFO] [logging.py:96:log_dist] [Rank 0] DeepSpeed info: version=0.15.1, git-hash=unknown, git-branch=unknown
[2025-06-30 01:11:14,851] [INFO] [config.py:733:__init__] Config mesh_device None world_size = 1
[2025-06-30 01:11:14,945] [INFO] [logging.py:96:log_dist] [Rank 0] DeepSpeed Flops Profiler Enabled: False
Using /home/xylv@iaclab.ic/.cache/torch_extensions/py38_cu121 as PyTorch extensions root...
Detected CUDA files, patching ldflags
Emitting ninja build file /home/xylv@iaclab.ic/.cache/torch_extensions/py38_cu121/fused_adam/build.ninja...
Building extension module fused_adam...
Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)
ninja: no work to do.
Loading extension module fused_adam...
Time to load fused_adam op: 0.16740107536315918 seconds
[2025-06-30 01:11:15,115] [INFO] [logging.py:96:log_dist] [Rank 0] Using DeepSpeed Optimizer param name adam as basic optimizer
[2025-06-30 01:11:15,115] [INFO] [logging.py:96:log_dist] [Rank 0] Removing param_group that has no 'params' in the basic Optimizer
[2025-06-30 01:11:15,116] [INFO] [logging.py:96:log_dist] [Rank 0] DeepSpeed Basic Optimizer = FusedAdam
[2025-06-30 01:11:15,116] [INFO] [utils.py:59:is_zero_supported_optimizer] Checking ZeRO support for optimizer=FusedAdam type=<class 'deepspeed.ops.adam.fused_adam.FusedAdam'>
[2025-06-30 01:11:15,116] [INFO] [logging.py:96:log_dist] [Rank 0] Creating torch.float32 ZeRO stage 2 optimizer
[2025-06-30 01:11:15,116] [INFO] [stage_1_and_2.py:148:__init__] Reduce bucket size 500000000
[2025-06-30 01:11:15,117] [INFO] [stage_1_and_2.py:149:__init__] Allgather bucket size 500000000
[2025-06-30 01:11:15,117] [INFO] [stage_1_and_2.py:150:__init__] CPU Offload: False
[2025-06-30 01:11:15,117] [INFO] [stage_1_and_2.py:151:__init__] Round robin gradient partitioning: False
[2025-06-30 01:11:15,338] [INFO] [utils.py:781:see_memory_usage] Before initializing optimizer states
[2025-06-30 01:11:15,339] [INFO] [utils.py:782:see_memory_usage] MA 0.01 GB         Max_MA 0.01 GB         CA 0.03 GB         Max_CA 0 GB 
[2025-06-30 01:11:15,339] [INFO] [utils.py:789:see_memory_usage] CPU Virtual Memory:  used = 13.56 GB, percent = 2.7%
[2025-06-30 01:11:15,527] [INFO] [utils.py:781:see_memory_usage] After initializing optimizer states
[2025-06-30 01:11:15,527] [INFO] [utils.py:782:see_memory_usage] MA 0.01 GB         Max_MA 0.01 GB         CA 0.03 GB         Max_CA 0 GB 
[2025-06-30 01:11:15,528] [INFO] [utils.py:789:see_memory_usage] CPU Virtual Memory:  used = 13.56 GB, percent = 2.7%
[2025-06-30 01:11:15,528] [INFO] [stage_1_and_2.py:543:__init__] optimizer state initialized
[2025-06-30 01:11:15,700] [INFO] [utils.py:781:see_memory_usage] After initializing ZeRO optimizer
[2025-06-30 01:11:15,701] [INFO] [utils.py:782:see_memory_usage] MA 0.01 GB         Max_MA 0.01 GB         CA 0.03 GB         Max_CA 0 GB 
[2025-06-30 01:11:15,701] [INFO] [utils.py:789:see_memory_usage] CPU Virtual Memory:  used = 13.56 GB, percent = 2.7%
[2025-06-30 01:11:15,702] [INFO] [logging.py:96:log_dist] [Rank 0] DeepSpeed Final Optimizer = DeepSpeedZeroOptimizer
[2025-06-30 01:11:15,702] [INFO] [logging.py:96:log_dist] [Rank 0] DeepSpeed using client callable to create LR scheduler
[2025-06-30 01:11:15,702] [INFO] [logging.py:96:log_dist] [Rank 0] DeepSpeed LR Scheduler = <torch.optim.lr_scheduler.OneCycleLR object at 0x7f42a0069f10>
[2025-06-30 01:11:15,702] [INFO] [logging.py:96:log_dist] [Rank 0] step=0, skipped=0, lr=[3.9999999999999996e-05], mom=['[0.9,0.99]']
[2025-06-30 01:11:15,702] [INFO] [config.py:999:print] DeepSpeedEngine configuration:
[2025-06-30 01:11:15,703] [INFO] [config.py:1003:print]   activation_checkpointing_config  {
    "partition_activations": true, 
    "contiguous_memory_optimization": true, 
    "cpu_checkpointing": false, 
    "number_checkpoints": null, 
    "synchronize_checkpoint_boundary": false, 
    "profile": false
}
[2025-06-30 01:11:15,703] [INFO] [config.py:1003:print]   aio_config ................... {'block_size': 1048576, 'queue_depth': 8, 'thread_count': 1, 'single_submit': False, 'overlap_events': True, 'use_gds': False}
[2025-06-30 01:11:15,703] [INFO] [config.py:1003:print]   amp_enabled .................. False
[2025-06-30 01:11:15,703] [INFO] [config.py:1003:print]   amp_params ................... False
[2025-06-30 01:11:15,703] [INFO] [config.py:1003:print]   autotuning_config ............ {
    "enabled": false, 
    "start_step": null, 
    "end_step": null, 
    "metric_path": null, 
    "arg_mappings": null, 
    "metric": "throughput", 
    "model_info": null, 
    "results_dir": "autotuning_results", 
    "exps_dir": "autotuning_exps", 
    "overwrite": true, 
    "fast": true, 
    "start_profile_step": 3, 
    "end_profile_step": 5, 
    "tuner_type": "gridsearch", 
    "tuner_early_stopping": 5, 
    "tuner_num_trials": 50, 
    "model_info_path": null, 
    "mp_size": 1, 
    "max_train_batch_size": null, 
    "min_train_batch_size": 1, 
    "max_train_micro_batch_size_per_gpu": 1.024000e+03, 
    "min_train_micro_batch_size_per_gpu": 1, 
    "num_tuning_micro_batch_sizes": 3
}
[2025-06-30 01:11:15,703] [INFO] [config.py:1003:print]   bfloat16_enabled ............. False
[2025-06-30 01:11:15,703] [INFO] [config.py:1003:print]   bfloat16_immediate_grad_update  False
[2025-06-30 01:11:15,703] [INFO] [config.py:1003:print]   checkpoint_parallel_write_pipeline  False
[2025-06-30 01:11:15,703] [INFO] [config.py:1003:print]   checkpoint_tag_validation_enabled  True
[2025-06-30 01:11:15,703] [INFO] [config.py:1003:print]   checkpoint_tag_validation_fail  False
[2025-06-30 01:11:15,703] [INFO] [config.py:1003:print]   comms_config ................. <deepspeed.comm.config.DeepSpeedCommsConfig object at 0x7f42a00697f0>
[2025-06-30 01:11:15,703] [INFO] [config.py:1003:print]   communication_data_type ...... None
[2025-06-30 01:11:15,703] [INFO] [config.py:1003:print]   compression_config ........... {'weight_quantization': {'shared_parameters': {'enabled': False, 'quantizer_kernel': False, 'schedule_offset': 0, 'quantize_groups': 1, 'quantize_verbose': False, 'quantization_type': 'symmetric', 'quantize_weight_in_forward': False, 'rounding': 'nearest', 'fp16_mixed_quantize': False, 'quantize_change_ratio': 0.001}, 'different_groups': {}}, 'activation_quantization': {'shared_parameters': {'enabled': False, 'quantization_type': 'symmetric', 'range_calibration': 'dynamic', 'schedule_offset': 1000}, 'different_groups': {}}, 'sparse_pruning': {'shared_parameters': {'enabled': False, 'method': 'l1', 'schedule_offset': 1000}, 'different_groups': {}}, 'row_pruning': {'shared_parameters': {'enabled': False, 'method': 'l1', 'schedule_offset': 1000}, 'different_groups': {}}, 'head_pruning': {'shared_parameters': {'enabled': False, 'method': 'topk', 'schedule_offset': 1000}, 'different_groups': {}}, 'channel_pruning': {'shared_parameters': {'enabled': False, 'method': 'l1', 'schedule_offset': 1000}, 'different_groups': {}}, 'layer_reduction': {'enabled': False}}
[2025-06-30 01:11:15,703] [INFO] [config.py:1003:print]   curriculum_enabled_legacy .... False
[2025-06-30 01:11:15,704] [INFO] [config.py:1003:print]   curriculum_params_legacy ..... False
[2025-06-30 01:11:15,704] [INFO] [config.py:1003:print]   data_efficiency_config ....... {'enabled': False, 'seed': 1234, 'data_sampling': {'enabled': False, 'num_epochs': 1000, 'num_workers': 0, 'curriculum_learning': {'enabled': False}}, 'data_routing': {'enabled': False, 'random_ltd': {'enabled': False, 'layer_token_lr_schedule': {'enabled': False}}}}
[2025-06-30 01:11:15,704] [INFO] [config.py:1003:print]   data_efficiency_enabled ...... False
[2025-06-30 01:11:15,704] [INFO] [config.py:1003:print]   dataloader_drop_last ......... False
[2025-06-30 01:11:15,704] [INFO] [config.py:1003:print]   disable_allgather ............ False
[2025-06-30 01:11:15,704] [INFO] [config.py:1003:print]   dump_state ................... False
[2025-06-30 01:11:15,704] [INFO] [config.py:1003:print]   dynamic_loss_scale_args ...... None
[2025-06-30 01:11:15,704] [INFO] [config.py:1003:print]   eigenvalue_enabled ........... False
[2025-06-30 01:11:15,704] [INFO] [config.py:1003:print]   eigenvalue_gas_boundary_resolution  1
[2025-06-30 01:11:15,704] [INFO] [config.py:1003:print]   eigenvalue_layer_name ........ bert.encoder.layer
[2025-06-30 01:11:15,704] [INFO] [config.py:1003:print]   eigenvalue_layer_num ......... 0
[2025-06-30 01:11:15,704] [INFO] [config.py:1003:print]   eigenvalue_max_iter .......... 100
[2025-06-30 01:11:15,704] [INFO] [config.py:1003:print]   eigenvalue_stability ......... 1e-06
[2025-06-30 01:11:15,704] [INFO] [config.py:1003:print]   eigenvalue_tol ............... 0.01
[2025-06-30 01:11:15,704] [INFO] [config.py:1003:print]   eigenvalue_verbose ........... False
[2025-06-30 01:11:15,704] [INFO] [config.py:1003:print]   elasticity_enabled ........... False
[2025-06-30 01:11:15,704] [INFO] [config.py:1003:print]   flops_profiler_config ........ {
    "enabled": false, 
    "recompute_fwd_factor": 0.0, 
    "profile_step": 1, 
    "module_depth": -1, 
    "top_modules": 1, 
    "detailed": true, 
    "output_file": null
}
[2025-06-30 01:11:15,704] [INFO] [config.py:1003:print]   fp16_auto_cast ............... None
[2025-06-30 01:11:15,704] [INFO] [config.py:1003:print]   fp16_enabled ................. False
[2025-06-30 01:11:15,704] [INFO] [config.py:1003:print]   fp16_master_weights_and_gradients  False
[2025-06-30 01:11:15,704] [INFO] [config.py:1003:print]   global_rank .................. 0
[2025-06-30 01:11:15,704] [INFO] [config.py:1003:print]   grad_accum_dtype ............. None
[2025-06-30 01:11:15,704] [INFO] [config.py:1003:print]   gradient_accumulation_steps .. 1
[2025-06-30 01:11:15,704] [INFO] [config.py:1003:print]   gradient_clipping ............ 1
[2025-06-30 01:11:15,704] [INFO] [config.py:1003:print]   gradient_predivide_factor .... 1.0
[2025-06-30 01:11:15,704] [INFO] [config.py:1003:print]   graph_harvesting ............. False
[2025-06-30 01:11:15,704] [INFO] [config.py:1003:print]   hybrid_engine ................ enabled=False max_out_tokens=512 inference_tp_size=1 release_inference_cache=False pin_parameters=True tp_gather_partition_size=8
[2025-06-30 01:11:15,704] [INFO] [config.py:1003:print]   initial_dynamic_scale ........ 65536
[2025-06-30 01:11:15,704] [INFO] [config.py:1003:print]   load_universal_checkpoint .... False
[2025-06-30 01:11:15,704] [INFO] [config.py:1003:print]   loss_scale ................... 0
[2025-06-30 01:11:15,704] [INFO] [config.py:1003:print]   memory_breakdown ............. False
[2025-06-30 01:11:15,704] [INFO] [config.py:1003:print]   mics_hierarchial_params_gather  False
[2025-06-30 01:11:15,704] [INFO] [config.py:1003:print]   mics_shard_size .............. -1
[2025-06-30 01:11:15,705] [INFO] [config.py:1003:print]   monitor_config ............... tensorboard=TensorBoardConfig(enabled=False, output_path='', job_name='DeepSpeedJobName') comet=CometConfig(enabled=False, samples_log_interval=100, project=None, workspace=None, api_key=None, experiment_name=None, experiment_key=None, online=None, mode=None) wandb=WandbConfig(enabled=False, group=None, team=None, project='deepspeed') csv_monitor=CSVConfig(enabled=False, output_path='', job_name='DeepSpeedJobName')
[2025-06-30 01:11:15,705] [INFO] [config.py:1003:print]   nebula_config ................ {
    "enabled": false, 
    "persistent_storage_path": null, 
    "persistent_time_interval": 100, 
    "num_of_version_in_retention": 2, 
    "enable_nebula_load": true, 
    "load_path": null
}
[2025-06-30 01:11:15,705] [INFO] [config.py:1003:print]   optimizer_legacy_fusion ...... False
[2025-06-30 01:11:15,705] [INFO] [config.py:1003:print]   optimizer_name ............... adam
[2025-06-30 01:11:15,705] [INFO] [config.py:1003:print]   optimizer_params ............. {'lr': 0.001, 'weight_decay': 0, 'bias_correction': True, 'eps': 1e-10, 'betas': '[0.9,0.99]'}
[2025-06-30 01:11:15,705] [INFO] [config.py:1003:print]   pipeline ..................... {'stages': 'auto', 'partition': 'best', 'seed_layers': False, 'activation_checkpoint_interval': 0, 'pipe_partitioned': True, 'grad_partitioned': True}
[2025-06-30 01:11:15,705] [INFO] [config.py:1003:print]   pld_enabled .................. False
[2025-06-30 01:11:15,705] [INFO] [config.py:1003:print]   pld_params ................... False
[2025-06-30 01:11:15,705] [INFO] [config.py:1003:print]   prescale_gradients ........... False
[2025-06-30 01:11:15,705] [INFO] [config.py:1003:print]   scheduler_name ............... None
[2025-06-30 01:11:15,705] [INFO] [config.py:1003:print]   scheduler_params ............. None
[2025-06-30 01:11:15,705] [INFO] [config.py:1003:print]   seq_parallel_communication_data_type  torch.float32
[2025-06-30 01:11:15,705] [INFO] [config.py:1003:print]   sparse_attention ............. None
[2025-06-30 01:11:15,705] [INFO] [config.py:1003:print]   sparse_gradients_enabled ..... False
[2025-06-30 01:11:15,705] [INFO] [config.py:1003:print]   steps_per_print .............. 10
[2025-06-30 01:11:15,705] [INFO] [config.py:1003:print]   timers_config ................ enabled=True synchronized=True
[2025-06-30 01:11:15,705] [INFO] [config.py:1003:print]   train_batch_size ............. 128
[2025-06-30 01:11:15,705] [INFO] [config.py:1003:print]   train_micro_batch_size_per_gpu  128
[2025-06-30 01:11:15,705] [INFO] [config.py:1003:print]   use_data_before_expert_parallel_  False
[2025-06-30 01:11:15,705] [INFO] [config.py:1003:print]   use_node_local_storage ....... False
[2025-06-30 01:11:15,705] [INFO] [config.py:1003:print]   wall_clock_breakdown ......... False
[2025-06-30 01:11:15,705] [INFO] [config.py:1003:print]   weight_quantization_config ... None
[2025-06-30 01:11:15,705] [INFO] [config.py:1003:print]   world_size ................... 1
[2025-06-30 01:11:15,705] [INFO] [config.py:1003:print]   zero_allow_untested_optimizer  False
[2025-06-30 01:11:15,705] [INFO] [config.py:1003:print]   zero_config .................. stage=2 contiguous_gradients=True reduce_scatter=True reduce_bucket_size=500000000 use_multi_rank_bucket_allreduce=True allgather_partitions=True allgather_bucket_size=500000000 overlap_comm=True load_from_fp32_weights=True elastic_checkpoint=False offload_param=None offload_optimizer=None sub_group_size=1000000000 cpu_offload_param=None cpu_offload_use_pin_memory=None cpu_offload=None prefetch_bucket_size=50000000 param_persistence_threshold=100000 model_persistence_threshold=9223372036854775807 max_live_parameters=1000000000 max_reuse_distance=1000000000 gather_16bit_weights_on_model_save=False use_all_reduce_for_fetch_params=False stage3_gather_fp16_weights_on_model_save=False ignore_unused_parameters=True legacy_stage1=False round_robin_gradients=False zero_hpz_partition_size=1 zero_quantized_weights=False zero_quantized_nontrainable_weights=False zero_quantized_gradients=False mics_shard_size=-1 mics_hierarchical_params_gather=False memory_efficient_linear=True pipeline_loading_checkpoint=False override_module_apply=True
[2025-06-30 01:11:15,705] [INFO] [config.py:1003:print]   zero_enabled ................. True
[2025-06-30 01:11:15,705] [INFO] [config.py:1003:print]   zero_force_ds_cpu_optimizer .. True
[2025-06-30 01:11:15,705] [INFO] [config.py:1003:print]   zero_optimization_stage ...... 2
[2025-06-30 01:11:15,706] [INFO] [config.py:989:print_user_config]   json = {
    "train_micro_batch_size_per_gpu": 128, 
    "fp16": {
        "enabled": false, 
        "loss_scale": 0, 
        "loss_scale_window": 1000, 
        "initial_scale_power": 16, 
        "hysteresis": 2, 
        "min_loss_scale": 1
    }, 
    "optimizer": {
        "type": "Adam", 
        "params": {
            "lr": 0.001, 
            "weight_decay": 0, 
            "bias_correction": true, 
            "eps": 1e-10, 
            "betas": "[0.9,0.99]"
        }
    }, 
    "zero_optimization": {
        "stage": 2, 
        "overlap_comm": true
    }, 
    "flops_profiler": {
        "enabled": false, 
        "profile_step": 1, 
        "module_depth": -1, 
        "top_modules": 1, 
        "detailed": true, 
        "output_file": null
    }, 
    "activation_checkpointing": {
        "partition_activations": true, 
        "contiguous_memory_optimization": true, 
        "cpu_checkpointing": false
    }, 
    "wall_clock_breakdown": false, 
    "gradient_accumulation_steps": 1, 
    "train_batch_size": 128, 
    "gradient_clipping": 1
}
[Debug] model-ema embedding_params:
tensor([[ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],
        [ 0.0098, -0.0145, -0.0088,  ...,  0.0123,  0.0107, -0.0084],
        [-0.0195, -0.0065,  0.0294,  ...,  0.0103, -0.0074, -0.0097],
        ...,
        [-0.0317,  0.0159, -0.0079,  ..., -0.0105,  0.0002, -0.0163],
        [ 0.0060, -0.0086,  0.0186,  ...,  0.0084,  0.0128, -0.0008],
        [-0.0044, -0.0031, -0.0171,  ...,  0.0191, -0.0359, -0.0188]],
       device='cuda:0')
[2025-06-30 01:11:15.818378] Finish -> Dump to `tokenization_config.json`
[2025-06-30 01:11:15.818524] Finish -> Dump to `params.txt`
[2025-06-30 01:11:15.818781] Finish -> Dump to `ds_config.json`
[2025-06-30 01:11:15.818802] Training start with j_init 0 and ep_init 0 ...
[2025-06-30 01:11:15.818812] running eval before training starts ...
[2025-06-30 01:11:18,925] [INFO] [real_accelerator.py:203:get_accelerator] Setting ds_accelerator to cuda (auto detect)
[2025-06-30 01:11:18,931] [INFO] [real_accelerator.py:203:get_accelerator] Setting ds_accelerator to cuda (auto detect)
[2025-06-30 01:11:19,180] [INFO] [real_accelerator.py:203:get_accelerator] Setting ds_accelerator to cuda (auto detect)
[2025-06-30 01:11:19,228] [INFO] [real_accelerator.py:203:get_accelerator] Setting ds_accelerator to cuda (auto detect)
[2025-06-30 01:11:19,241] [INFO] [real_accelerator.py:203:get_accelerator] Setting ds_accelerator to cuda (auto detect)
[2025-06-30 01:11:19,280] [INFO] [real_accelerator.py:203:get_accelerator] Setting ds_accelerator to cuda (auto detect)
[2025-06-30 01:11:19,299] [INFO] [real_accelerator.py:203:get_accelerator] Setting ds_accelerator to cuda (auto detect)
[2025-06-30 01:11:19,309] [INFO] [real_accelerator.py:203:get_accelerator] Setting ds_accelerator to cuda (auto detect)
[2025-06-30 01:11:19,365] [INFO] [real_accelerator.py:203:get_accelerator] Setting ds_accelerator to cuda (auto detect)
[2025-06-30 01:11:19,406] [INFO] [real_accelerator.py:203:get_accelerator] Setting ds_accelerator to cuda (auto detect)
[DataLoader Worker 1] seed `random` & `np.random` with 3755130380448426379 & 2078796171!
[DataLoader Worker 3] seed `random` & `np.random` with 3755130380448426381 & 2078796173!
[DataLoader Worker 7] seed `random` & `np.random` with 3755130380448426385 & 2078796177!
[DataLoader Worker 6] seed `random` & `np.random` with 3755130380448426384 & 2078796176!
[DataLoader Worker 9] seed `random` & `np.random` with 3755130380448426387 & 2078796179!
[DataLoader Worker 0] seed `random` & `np.random` with 3755130380448426378 & 2078796170!
Traceback (most recent call last):
  File "./examples/train_supervised.py", line 845, in <module>
    fire.Fire(train)
  File "/home/xylv@iaclab.ic/anaconda3/envs/graphgpt/lib/python3.8/site-packages/fire/core.py", line 141, in Fire
    component_trace = _Fire(component, args, parsed_flag_args, context, name)
  File "/home/xylv@iaclab.ic/anaconda3/envs/graphgpt/lib/python3.8/site-packages/fire/core.py", line 475, in _Fire
    component, remaining_args = _CallAndUpdateTrace(
  File "/home/xylv@iaclab.ic/anaconda3/envs/graphgpt/lib/python3.8/site-packages/fire/core.py", line 691, in _CallAndUpdateTrace
    component = fn(*varargs, **kwargs)
  File "./examples/train_supervised.py", line 523, in train
    val_loss, val_cls_metrics, val_ogb_eval_res, val_triplet = evaluate(
  File "/home/xylv@iaclab.ic/anaconda3/envs/graphgpt/lib/python3.8/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "./examples/train_supervised.py", line 78, in evaluate
    for j, test_data in enumerate(loader, 1):
  File "/home/xylv@iaclab.ic/anaconda3/envs/graphgpt/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 630, in __next__
    data = self._next_data()
  File "/home/xylv@iaclab.ic/anaconda3/envs/graphgpt/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 1345, in _next_data
    return self._process_data(data)
  File "/home/xylv@iaclab.ic/anaconda3/envs/graphgpt/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 1371, in _process_data
    data.reraise()
  File "/home/xylv@iaclab.ic/anaconda3/envs/graphgpt/lib/python3.8/site-packages/torch/_utils.py", line 694, in reraise
    raise exception
IndexError: Caught IndexError in DataLoader worker process 0.
Original Traceback (most recent call last):
  File "/home/xylv@iaclab.ic/anaconda3/envs/graphgpt/lib/python3.8/site-packages/torch/utils/data/_utils/worker.py", line 308, in _worker_loop
    data = fetcher.fetch(index)
  File "/home/xylv@iaclab.ic/anaconda3/envs/graphgpt/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py", line 54, in fetch
    return self.collate_fn(data)
  File "/home/xylv@iaclab.ic/python_code/analog-gpt/./src/data/collator.py", line 91, in __call__
    [
  File "/home/xylv@iaclab.ic/python_code/analog-gpt/./src/data/collator.py", line 92, in <listcomp>
    _add_idx_to_dict(self.tokenizer(graph, self.is_training), idx)
  File "/home/xylv@iaclab.ic/python_code/analog-gpt/./src/data/tokenizer.py", line 3870, in __call__
    token_res = self.tokenize(graph)
  File "/home/xylv@iaclab.ic/python_code/analog-gpt/./src/data/tokenizer.py", line 3678, in tokenize
    ) = get_semantics_raw_node_edge2attr_mapping(path, graph, self.config)
  File "/home/xylv@iaclab.ic/python_code/analog-gpt/./src/data/tokenizer.py", line 3193, in get_semantics_raw_node_edge2attr_mapping
    dict_map_node = get_semantics_attr_mapping(
  File "/home/xylv@iaclab.ic/python_code/analog-gpt/./src/data/tokenizer.py", line 3142, in get_semantics_attr_mapping
    tmp_map = func_attr_mapping(path, data, discrete_attr)
  File "/home/xylv@iaclab.ic/python_code/analog-gpt/./src/data/tokenizer.py", line 3050, in _get_node2attr_mapping
    tmp_map[src] = data[attr_name][src].numpy()
IndexError: index 279365 is out of bounds for dimension 0 with size 117

[zhaoqf123]

I will take a look at the example dataset and get back to you soon.

[xylyu-zju]

Thank you, I just found that when I switched the edge index from global index to the local index, the error was solved. But that's not what I want. I want to us so muche the global node and edge idx to train the model because my data sets have mutual inclusion relationships, for example, data set 4 may appear as part of data set 1000. I want to know how to use the global index for graph classification in my dataset.

[zhaoqf123]

Thank you, I just found that when I switched the edge index from global index to the local index, the error was solved. But that's not what I want. I want to us so muche the global node and edge idx to train the model because my data sets have mutual inclusion relationships, for example, data set 4 may appear as part of data set 1000. I want to know how to use the global index for graph classification in my dataset.

Since each graph, i.e., analog circuit, is independent, no need to explicitly to indicate that they share the same small components. So we use local index.

For example, in NLP, two independent sentences may both have the same short phrase "take a look at", but we don't need explicitly state that the phrase are the same. Let the model resolve this by itself.

[zhaoqf123]

The script for processing analog circuit data to torch-geometric format, and the config to train such dataset have been updated in this commit 856819e, after some discussion and collaboration with @xylyu-zju. Thanks!