alibaba / alibaba/graphlearn-for-pytorch

bad training performance with intel CPU

Open
#165 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
146
Forks
48
PR merge metrics
No merged PRs in 30d

Description

### 🐛 Describe the bug

I tried to test pure local CPU train speed on some PyG datasets like Reddit on an intel CPU server, but when it comes to Reddit2 dataset, the performance degraded significantly. One epoch needed to cost more than 700s to finish, while PyG only cost around 80s. Then I turn to an AMD CPU server and tested again, everything worked fine. So is there any problem with intel CPUs? Here is my test code:
```python
import time
import torch
import torch.nn as nn
import graphlearn_torch as glt

from torch_geometric.datasets import Reddit, Reddit2, Flickr, AmazonProducts
from torch_geometric.nn import SAGEConv

device = torch.device("cpu")
t_init_start = time.perf_counter()

dataset = Reddit2()
data = dataset[0]

x = data.x.cpu()
y = data.y.cpu()
edge_index = data.edge_index.cpu()
train_mask = data.train_mask

num_nodes = data.num_nodes
num_classes = dataset.num_classes

print(f"Reddit loaded: nodes={num_nodes}, edges={edge_index.size(1)}")

t_init_end = time.perf_counter()

t_graph_start = time.perf_counter()

glt_dataset = glt.data.Dataset()

glt_dataset.init_graph(
edge_index=edge_index,
graph_mode='cpu'
)

glt_dataset.init_node_features(
node_feature_data=x,
with_gpu=False
)

glt_dataset.init_node_labels(node_label_data=y)

t_graph_end = time.perf_counter()
print(f"GraphStore construction time: {t_graph_end - t_graph_start:.2f}s")

class GraphSAGE(nn.Module):
def __init__(self, in_dim, hidden_dim, out_dim):
super().__init__()
self.conv1 = SAGEConv(in_dim, hidden_dim)
self.conv2 = SAGEConv(hidden_dim, out_dim)
self.relu = nn.ReLU()

def forward(self, x, edge_index):
x = self.relu(self.conv1(x, edge_index))
x = self.conv2(x, edge_index)
return x

model = GraphSAGE(
in_dim=x.size(1),
hidden_dim=256,
out_dim=num_classes
).to(device)

optimizer = torch.optim.Adam(model.parameters(), lr=0.003, weight_decay=5e-4)
loss_fn = nn.CrossEntropyLoss()

train_nids = train_mask.nonzero(as_tuple=True)[0]
train_labels = y[train_nids]

train_loader = glt.loader.NeighborLoader(
data=glt_dataset,
input_nodes=train_nids,
num_neighbors=[15, 10],
batch_size=64,
shuffle=True,
device=torch.device("cpu"),
)

num_epochs = 1
t_train_start = time.perf_counter()
step = 0

for epoch in range(num_epochs):
model.train()
epoch_loss = 0.0
epoch_samples = 0
t_epoch_start = time.perf_counter()

for batch in train_loader:
x_batch = batch.x.to(device)
y_batch = batch.y.to(device)
edge_index_batch = batch.edge_index.to(device)

logits = model(x_batch, edge_index_batch)

loss = loss_fn(logits[:batch.batch_size], y_batch[:batch.batch_size])

optimizer.zero_grad()
loss.backward()
optimizer.step()

epoch_loss += loss.item() * batch.batch_size
epoch_samples += batch.batch_size

step += 1
if step % 50 == 0:
print(f'Step {step}, loss={loss.item():.4f}')

t_epoch_end = time.perf_counter()
print(f"Epoch {epoch:02d} | Loss {epoch_loss / epoch_samples:.4f} | Time {t_epoch_end - t_epoch_start:.2f}s")

t_train_end = time.perf_counter()

print("\n===== Timing Summary (GLT, CPU) =====")
print(f"Initialization Time : {t_init_end - t_init_start:.2f}s")
print(f"GraphStore Build : {t_graph_end - t_graph_start:.2f}s")
print(f"Training Time : {t_train_end - t_train_start:.2f}s")
```

### Environment

* GLT version:0.2.4
* PyG version:2.7.0
* PyTorch version:2.6.0
* OS:Ubuntu 24.04.3 LTS
* Python version:3.10
* CUDA/cuDNN version:12.8
* Any other relevant information
* CPU: intel 5218R*2

Contributor guide

No contributing guide indexed for this repository

Research direction

No source file or failing test is named. Start by running the supplied Reddit2 GraphSAGE and NeighborLoader script with the stated versions on the Intel 5218R and AMD systems, then compare the initialization, graph construction, and training timings. Done means identifying and documenting the Intel-specific cause or confirming the relevant performance limitation.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.