Let’s evaluate the tinytopics topic model training speed on CPU vs. GPU +on mainstream consumer hardware using simulated data. We will compare +the time consumed under combinations of the three key parameters +defining the problem size:
+n).m).k).Experiment environment:
+n) grows,
+ on both CPU and GPU.k) grows.n and k fixed and vocabulary size (m) grows, CPU time will
+ grow linearly while GPU time stays constant. For m larger than a
+ certain threshold, training on GPU will be faster than CPU.import time
+
+import torch
+import pandas as pd
+import matplotlib.pyplot as plt
+import tinytopics as tt
+Set seed for reproducibility:
+ +Define parameter grids:
+n_values = [1000, 5000] # Number of documents
+m_values = [1000, 5000, 10000, 20000] # Vocabulary size
+k_values = [10, 50, 100] # Number of topics
+avg_doc_length = 256 * 256
+Create a data frame to store the benchmark results.
+benchmark_results = pd.DataFrame()
+
+def benchmark(X, k, device):
+ start_time = time.time()
+ model, losses = tt.fit_model(X, k, device=device)
+ elapsed_time = time.time() - start_time
+
+ return elapsed_time
+for n in n_values:
+ for m in m_values:
+ for k in k_values:
+ print(f"Benchmarking for n={n}, m={m}, k={k}...")
+
+ X, true_L, true_F = tt.generate_synthetic_data(n, m, k, avg_doc_length=avg_doc_length)
+
+ # Benchmark on CPU
+ cpu_time = benchmark(X, k, torch.device("cpu"))
+ cpu_result = pd.DataFrame([{"n": n, "m": m, "k": k, "device": "CPU", "time": cpu_time}])
+
+ if not cpu_result.isna().all().any():
+ benchmark_results = pd.concat([benchmark_results, cpu_result], ignore_index=True)
+
+ # Benchmark on GPU if available
+ if torch.cuda.is_available():
+ gpu_time = benchmark(X, k, torch.device("cuda"))
+ gpu_result = pd.DataFrame([{"n": n, "m": m, "k": k, "device": "GPU", "time": gpu_time}])
+
+ if not gpu_result.isna().all().any():
+ benchmark_results = pd.concat([benchmark_results, gpu_result], ignore_index=True)
+Save results to a CSV file:
+ +Plot the number of terms (m) against the time consumed, conditioning
+on the number of documents (n), for each number of topics (k).
unique_series = len(n_values) * (2 if torch.cuda.is_available() else 1)
+colormap = tt.scale_color_tinytopics(unique_series)
+colors_list = [colormap(i) for i in range(unique_series)]
+
+for k in k_values:
+ plt.figure(figsize=(7, 4.3), dpi=300)
+
+ color_idx = 0
+ for n in n_values:
+ subset = benchmark_results[
+ (benchmark_results["n"] == n) & (benchmark_results["k"] == k)
+ ]
+
+ # Plot CPU results with a specific color
+ plt.plot(
+ subset[subset["device"] == "CPU"]["m"],
+ subset[subset["device"] == "CPU"]["time"],
+ label=f"CPU (n={n})",
+ linestyle="--",
+ marker="o",
+ color=colors_list[color_idx],
+ )
+ color_idx += 1
+
+ # Plot GPU results if available
+ if torch.cuda.is_available():
+ plt.plot(
+ subset[subset["device"] == "GPU"]["m"],
+ subset[subset["device"] == "GPU"]["time"],
+ label=f"GPU (n={n})",
+ linestyle="-",
+ marker="x",
+ color=colors_list[color_idx],
+ )
+ color_idx += 1
+
+ plt.xlabel("Vocabulary size (m)")
+ plt.ylabel("Training time (seconds)")
+ plt.title(f"Training time vs. vocabulary size (k={k})")
+ plt.legend()
+ plt.grid(True)
+ plt.savefig(f"training-time-k-{k}.png", dpi=300)
+ plt.close()
+
Tip
+The code from this article is available in:
+ +Follow the instructions in the article to run the example.
+tinytopics >= 0.7.0 supports distributed training using Hugging Face +Accelerate. This article +demonstrates how to run distributed training on a single node with +multiple GPUs.
+The example utilizes Distributed Data Parallel (DDP) for distributed +training. This approach assumes that the model parameters fit within the +memory of a single GPU, as each GPU maintains a synchronized copy of the +model. The input data can exceed the memory capacity. This is generally +a reasonable assumption for topic modeling tasks, as storing the +factorized matrices is often less memory-intensive.
+Hugging Face Accelerate also supports other distributed training +strategies such as Fully Sharded Data Parallel (FSDP) and DeepSpeed, +which distribute model tensors across different GPUs and allow training +larger models at the cost of speed.
+We will use a 100k x 100k count matrix with 20 topics for distributed
+training. To generate the example data, save the following code to
+distributed_data.py and run:
import os
+
+import numpy as np
+
+import tinytopics as tt
+
+
+def main():
+ n, m, k = 100_000, 100_000, 20
+ data_path = "X.npy"
+
+ if os.path.exists(data_path):
+ print(f"Data already exists at {data_path}")
+ return
+
+ print("Generating synthetic data...")
+ tt.set_random_seed(42)
+ X, true_L, true_F = tt.generate_synthetic_data(
+ n=n, m=m, k=k, avg_doc_length=256 * 256
+ )
+
+ print(f"Saving data to {data_path}")
+ X_numpy = X.cpu().numpy()
+ np.save(data_path, X_numpy)
+
+
+if __name__ == "__main__":
+ main()
+Generating the data is time-consuming (about 10 minutes), so running it +as a standalone script helps avoid potential timeout errors during +distributed training. You can also execute it on an instance type +suitable for your data ingestion pipeline, rather than using valuable +GPU instance hours.
+First, configure the distributed environment by running:
+ +You will be prompted to answer questions about the distributed training +environment and strategy. The answers will be saved to a configuration +file at:
+~/.cache/huggingface/accelerate/default_config.yaml
+You can rerun accelerate config at any time to update the
+configuration. For data distributed parallel on a 4-GPU node, select
+single-node multi-GPU training options with the number of GPUs set to 4,
+and use the default settings for the remaining questions (mostly “no”).
Next, save the following code to distributed_training.py and run:
import os
+
+from accelerate import Accelerator
+from accelerate.utils import set_seed
+
+import tinytopics as tt
+
+
+def main():
+ accelerator = Accelerator()
+ set_seed(42)
+ k = 20
+ data_path = "X.npy"
+
+ if not os.path.exists(data_path):
+ raise FileNotFoundError(
+ f"{data_path} not found. Run distributed_data.py first."
+ )
+
+ print(f"Loading data from {data_path}")
+ X = tt.NumpyDiskDataset(data_path)
+
+ # All processes should have the data before proceeding
+ accelerator.wait_for_everyone()
+
+ model, losses = tt.fit_model_distributed(X, k=k)
+
+ # Only the main process should plot the loss
+ if accelerator.is_main_process:
+ tt.plot_loss(losses, output_file="loss.png")
+
+
+if __name__ == "__main__":
+ main()
+This script uses fit_model_distributed() (added in tinytopics 0.7.0)
+to train the model. Since distributed training on large datasets likely
+takes longer, fit_model_distributed() displays more detailed progress
+bars for each epoch, going through all batches in each epoch.
We ran the distributed training example on a 1-GPU, 4-GPU, and 8-GPU +setup with H100 and A100 GPUs. The table below shows the training time +per epoch, total time, GPU utilization, VRAM usage, instance cost, and +total cost.
+| Metric | +1x H100 (80 GB SXM5) | +4x H100 (80 GB SXM5) | +8x A100 (40 GB SXM4) | +
|---|---|---|---|
| Time per epoch (s) | +24 | +6 | +6 | +
| Total time (min) | +80 | +20 | +20 | +
| GPU utilization | +16% | +30-40% | +30-50% | +
| VRAM usage | +1% | +4% | +4% | +
| Instance cost (USD/h) | +3.29 | +12.36 | +10.32 | +
| Total cost (USD) | +4.38 | +4.12 | +3.44 | +
Using 4 H100 GPUs is approximately 4x faster than using 1 H100 GPU, with +a slightly lower total cost. Using 8x A100 GPUs has similar speed +comparing to 4x H100 GPUs but with an even lower total cost due to the +lower instance cost.
+The loss plot and real-time GPU utilization monitoring via nvtop on
+the 4x H100 GPU instance are shown below.
For more technical details on distributed training, please refer to the +Hugging Face Accelerate documentation, as this article covers only the +basics.
+ + + + + + + + + + + + + +Let’s walk through a canonical tinytopics workflow using a synthetic +dataset.
+Set random seed for reproducibility:
+ +Generate a synthetic dataset:
+n, m, k = 5000, 1000, 10
+X, true_L, true_F = tt.generate_synthetic_data(n, m, k, avg_doc_length=256 * 256)
+Fit the topic model and plot the loss curve. There will be a progress +bar.
+ +
Tip
+By default, tinytopics uses AdamW with weight decay as the optimizer, +and the cosine annealing with warm restarts scheduler. +This combination should help reduce the need of extensive manual tuning +of hyperparameters such as the learning rate. For optimal performance, +exploring the possible tuning parameter space is still recommended.
+Get the learned L and F matrices from the fitted topic model:
+ +To make it easier to inspect the results visually, we should try to +“align” the learned topics with the ground truth topics by their terms +similarity.
+aligned_indices = tt.align_topics(true_F, learned_F)
+learned_F_aligned = learned_F[aligned_indices]
+learned_L_aligned = learned_L[:, aligned_indices]
+Sort the documents in both the true document-topic matrix and the +learned document-topic matrix, grouped by dominant topics.
+sorted_indices = tt.sort_documents(true_L)
+true_L_sorted = true_L[sorted_indices]
+learned_L_sorted = learned_L_aligned[sorted_indices]
+Note
+The alignment step mostly only applies to simulation studies +because we often don't know the ground truth L and F for real datasets.
+We can use a “Structure plot” to visualize and compare the +document-topic distributions.
+tt.plot_structure(
+ true_L_sorted,
+ normalize_rows=True,
+ title="True document-topic distributions (sorted)",
+ output_file="L-true.png",
+)
+
tt.plot_structure(
+ learned_L_sorted,
+ normalize_rows=True,
+ title="Learned document-topic distributions (sorted and aligned)",
+ output_file="L-learned.png",
+)
+
We can also plot the top terms for each topic using bar charts.
+tt.plot_top_terms(
+ true_F,
+ n_top_terms=15,
+ title="Top terms per topic - true F matrix",
+ output_file="F-top-terms-true.png",
+)
+
tt.plot_top_terms(
+ learned_F_aligned,
+ n_top_terms=15,
+ title="Top terms per topic - learned F matrix (aligned)",
+ output_file="F-top-terms-learned.png",
+)
+