# Trainy (full content) > Trainy builds GPU infrastructure for foundation-model and research teams. It ships two products: Konduktor (GPU cluster scheduling and orchestration on Kubernetes) and Pluto (experiment tracking). Both products expose an MCP server, so an AI agent such as Claude can launch and restart Konduktor jobs, query metrics and logs, and analyze Pluto run data directly. ## Konduktor (https://www.trainy.ai/training) Konduktor is GPU cluster scheduling and orchestration on Kubernetes. Teams submit batch jobs with a simple YAML file and one CLI command (konduktor launch job.yaml), with no training-code changes. It provides priority queueing with preemption (high-priority jobs pause lower ones and resume them on completion), multi-node distributed training over InfiniBand and RoCE, continuous GPU and node health checks with automatic cordon and reschedule onto healthy hardware, and Prometheus/Grafana observability for GPU, network, and cluster telemetry. Konduktor exposes jobs over the Model Context Protocol, so Claude can launch workloads, read job and cluster logs, pull GPU metrics, and restart jobs with user approval. Pricing is quote based; book a demo at https://www.trainy.ai/book-a-demo. Docs: https://docs.trainy.ai/overview ## Inference (https://www.trainy.ai/inference) Konduktor schedules training and inference on the same GPUs: inference scales with traffic and training fills the gaps, so models are served from the same cluster they are trained on, with autoscaling including scale to zero. ## Pluto (https://www.trainy.ai/pluto) Pluto is open source experiment tracking for metrics, logs, files, and run forking, with a UI that stays responsive at thousands of runs. It is a drop-in replacement for Weights & Biases and Neptune: compatibility layers allow dual-logging to both platforms simultaneously so teams can validate Pluto against production usage before switching. Run forking branches a new experiment from any step of an existing run (pluto.init(fork_run_id=..., fork_step=..., config={...})). Pluto's MCP integration gives AI assistants read/query access to runs: list_projects, list_runs, get_run, query_logs, query_metrics, get_files, get_statistics (statistics and anomaly detection on metrics), and compare_runs. Connect with: claude mcp add pluto --transport http --header "Authorization: Bearer mlpi_xxxx" -- https://pluto-mcp.trainy.ai/mcp/ Pricing: free hosted tier with 2 seats and 2 GB storage; hosted Pro is $250 per seat per month with up to 10 seats, 10 billion data points per month, 10 TB storage, unlimited tracked hours, and run forking; Enterprise (more than 10 seats, SSO, audit logs, single-tenant or on-prem) is contact-us. Also free to self-host (open source, full source access, no usage limits), and academics get hosted Pluto free by signing up with a .edu email. Docs: https://docs.trainy.ai/pluto/overview ## Pricing (https://www.trainy.ai/pricing) Konduktor: custom quote sized to your cluster (BYOC platform fee plus per-node pricing). Pluto: free hosted tier (2 seats, 2 GB storage), $250 per seat per month hosted Pro (up to 10 seats), Enterprise for more than 10 seats, free academic tier (.edu email), and free self-host. YC companies get startup perks (free performance-tuning consultation, Konduktor discounts, 1 year of free Pluto Pro). ## Company Trainy is based in San Mateo, California (1900 S Norfolk St #350, San Mateo, CA 94403). Contact: sales@trainy.ai. GitHub: https://github.com/Trainy-ai. Blog feed: https://www.trainy.ai/feed.xml --- ## Blog: What Neptune Got Right (And How to Keep It) (https://www.trainy.ai/blog/what-neptune-got-right) Published 2026-02-18 by Roanak Baviskar. Neptune is shutting down on March 5, 2026 (10am PST). For teams that have been using it for years, the real loss isn't just the migration work. It's losing one of the few trackers that stayed fast when projects got big. At [Trainy](https://www.trainy.ai/), we build infrastructure for GPU clusters. When the shutdown was announced, our customers started asking what they should migrate to. The recurring requirement wasn't "more features", it was responsiveness with thousands of runs and heavy time-series logging. That's what Neptune consistently nailed, and none of the alternatives our customers evaluated matched it. So we built [Pluto](https://pluto.trainy.ai/): an experiment tracker based on [our fork](https://github.com/Trainy-ai/pluto) of the open source project [MLOp](https://github.com/mlop-ai/mlop). Same focus on staying snappy at scale, with a low-risk migration path. Dual-log by adding one import, validate parity, then cut over when you're ready. This post is about what made Neptune good, what we think could be better, and how we approached building Pluto. ## What Neptune Got Right ### Scalability and responsiveness Neptune handled scale well. You could log 50k+ metrics per run with no dropped points and no timeouts under sustained logging. Projects with thousands of runs remained usable. And the UI was fast: graphs loaded quickly, filtering the runs table didn't freeze the page. This sounds basic, but if you've used a slow experiment tracker while trying to debug a training run, you know how much it matters. ### The query system Neptune's query language (NQL) was actually quite powerful. You could filter runs by any field: `recall > 0.9 AND learning_rate <= 0.005`. You could combine conditions with AND/OR, use comparison operators on metrics, filter by tags, owners, dates. The query builder made it easy to construct these without memorizing syntax. ### Side-by-side comparison Neptune's side-by-side view was well designed. It stacked attributes as rows instead of columns, which made scrolling through differences easier. You could filter to show only rows with differences, set a reference run to see relative changes, and even diff long strings character-by-character. ## Where We See Room to Build ### Graph UX Neptune's graphing was solid, but there were a few areas where we think the experience could be better: - Legend sorting was based on current metric value, which meant the order shifted as you scrubbed through time - Filtering views within graphs required manually hiding metrics one by one - Smoothing was global rather than per-panel or per-metric - Syncing axes across graphs for debugging gradient spikes required manual axis linking per chart group ### Dev Workflow Neptune lived somewhat separately from the rest of your workflow. We'd want the ability to tie experiments to Linear/Jira issues, terminate running jobs from the UI, and view code diffs in the UI. That last one came up a lot for us: two runs, identical hyperparameters, different results. Something in the code changed, but you'd have to track down git commits and diff them yourself. ### Tensor logging Neptune required you to render plots client-side before logging them. There was no way to natively plot raw tensors on the fly. Logging tensors directly and generating visualizations server-side would let teams iterate on what they visualize without worrying about rendering and saving images client-side. ### LLM Integration Experiment data is a natural fit for LLM querying. Being able to ask "which runs last week had the lowest validation loss" or get automatic summaries of experiment batches would save a lot of manual dashboard digging. (We shipped an early version of this with [Pluto MCP](https://docs.trainy.ai/pluto/mcp), currently in alpha.) ## Pluto's Stack Pluto is built with the same focus on responsiveness that made Neptune good. It's based on [our fork of MLOp](https://github.com/Trainy-ai/pluto), and we've been working with design partners to figure out what actually matters to teams running serious training workloads. Under the hood: - ClickHouse for OLAP, enabling fast analytics on metric series with 1M+ points. - Postgres OLTP for user, organization, and run metadata. - Rust ingestion server for fast non-blocking metric writes. We added on-disk buffering to both the client and server for increased reliability against networking hiccups. - Ingestion server handles 2000+ req/s end-to-end per replica. P95 latency: 200ms at 2k req/s. - React frontend for a responsive UX. We added lazy loading optimizations to increase the responsiveness of the UI even when multiple runs and charts are being loaded on a single page. - Whole stack is self-hostable via docker-compose in about a minute. While we host this stack with multiple replicas for high availability, it's still very fast on a single machine with around 8 CPUs, 32GB RAM, and 1TB disk. {% youtube-video id="fjpRP0KLCug" /%} ### Neptune compatibility layer We built a compatibility layer so you can dual-log to both Neptune and Pluto with minimal code changes: ```python import pluto.compat.neptune # Add this line from neptune_scale import Run # Your existing Neptune Scale import # Your existing Neptune Scale code works unchanged run = Run(experiment_name="my-experiment") run.log_configs({"lr": 0.001, "batch_size": 32}) run.log_metrics({"train/loss": 0.5}, step=0) run.close() ``` Set the `PLUTO_PROJECT` environment variable and your metrics log to both platforms. This is intentional. We wanted to eliminate all risk of trying Pluto. Your Neptune logging keeps working exactly as before. If Pluto breaks, you still have Neptune. If you don't like Pluto, remove the import and nothing changes. There's no commitment until you're ready. When you are ready to stop logging to Neptune entirely: ```bash export DISABLE_NEPTUNE_LOGGING=true ``` All Neptune API calls get redirected to Pluto. No code changes required. ## Migration Guide ### Common pitfalls Before diving in, a few things worth knowing: - **Exports take longer than you'd expect.** Neptune explicitly warns that workspace exports can take days to weeks depending on size. Don't wait until the last week of February. - **Artifact sizes and storage planning.** If you've been logging model checkpoints, large files, or images, figure out how much storage you'll need on the other side before you start. - **Legacy Neptune client vs Neptune Scale.** Pluto's compatibility layer targets [Neptune Scale](https://github.com/neptune-ai/neptune-client-scale) (`neptune-scale` package), not the legacy `neptune` client. If your codebase still uses the old client, you'll need to [migrate to Neptune Scale](https://docs.neptune.ai/migration_neptune/) first. - **Run IDs and fork mapping.** Neptune's internal run IDs and fork relationships don't have a universal standard. Make sure your export preserves the lineage you care about. - **Verifying parity means more than metrics.** When we say "validate that everything matches," that includes step counts, tags, system metrics, and logged files, not just the loss curves. ### Step 1: Set up dual-logging Add the compatibility import and set the `PLUTO_PROJECT` environment variable as described in the compatibility layer section above. Run this on real training workloads to validate at scale. ### Step 2: Export your Neptune history We have an exporter that brings your historical runs into Pluto. Details: [Exporting Neptune Runs](https://docs.trainy.ai/pluto/exporting-neptune-runs) ### Step 3: Validate and give us feedback Before removing Neptune: - Do metrics match between both platforms? - Is the UI responsive at your scale? - Can your team find what they need? This is where we need your help. If something's missing, broken, or would block your migration, tell us. Email roanak@trainy.ai or open an issue on GitHub. We're shipping fixes fast and your feedback directly shapes what we build next. ### Step 4: Cut over Set `DISABLE_NEPTUNE_LOGGING=true`. All Neptune calls redirect to Pluto. Done. ## Pricing and Links Pluto's hosted plan is $250/seat/month, matching Neptune's pricing. - [Live playground](https://demo.pluto.trainy.ai/o/dev-org/projects/my-ml-project), no signup required - [Quickstart](https://docs.trainy.ai/pluto/quickstart) - [Neptune migration docs](https://docs.trainy.ai/pluto/neptune-migration) - [GitHub](https://github.com/Trainy-ai/pluto) We're listed on [Neptune's official transition hub](https://docs.neptune.ai/transition_hub/migration/to_pluto). Pluto is open-source. The [Python client](https://github.com/Trainy-ai/pluto) is Apache-2.0 licensed. The [server](https://github.com/Trainy-ai/pluto-server) is AGPL-3.0. You can self-host the entire stack via docker-compose: clone the server repo, copy `.env.example` to `.env`, and run `docker compose up --build`. The server README has full setup instructions. Questions? Email roanak@trainy.ai or [book a time](https://calendly.com/roanak/trainy-demo). --- ## Blog: Preflight Checks for GPU Fabric before Multi-node Training (https://www.trainy.ai/blog/preflight-checks-gpu-fabric) Published 2025-03-17 by Grace Boyle. Given the scale of modern LLMs, having a performant GPU network fabric is one of the most important factor for your training speed. The ability to efficiently communicate across hundreds of nodes allows you to maximize GPU utilization and explore more complex training parallelisms such as in [**Deepspeed**](https://www.deepspeed.ai/) or [**Fully Sharded Data Parallel**](https://pytorch.org/docs/stable/fsdp.html). These critically rely on high performance networking solutions requiring specialized hardware and software like InfiniBand or RDMA over Converged Ethernet (RoCE) with current vendor options often going up to now 3200 Gbps. ## Estimating Step Time for GPT language models To illustrate, how important GPU fabric is to training, let's consider training an 7B GPT-style language model trained at 2048 context length using a pair of H100x8 nodes and performing an all-reduce operation which is used to average gradients over GPUs. Typical model flops utilization (MFU) for falls [**around ~40%**](https://github.com/mosaicml/llm-foundry/tree/main/scripts/train/benchmarking). For an H100, that roughly translates to `396 TFLOP/second`. To perform, forward and backward pass of the model with reactivation using BF16, the [**number of TFLOP can be estimated as**](https://github.com/bigscience-workshop/bigscience/tree/master/math#calculate-tflops): ``` tflop = model_size_in_B * 4 * 2 * seqlen * global_batch_size / (total_gpus * 1e3) ``` We already have `seqlen = 2048` and `model_size_in_B = 7`. A [**typical batch size**](https://github.com/mosaicml/llm-foundry/tree/main/scripts/train/benchmarking) might be `global_batch_size = 16`. This gives us 114 TFLOP. So we expect, computing each step to take roughly `114 TFLOP / 396 TFLOP/second = 0.29 seconds`. Now we need to estimate the time it takes to average the gradients over all the GPUs. For a 7B model, we'll pass gradients in fp32 meaning we need to pass S := 7 Billion x 4 Bytes = 28 GB in gradients every iteration. With n GPUs, for all-reduce, each GPU communicates its gradients to all the others resulting in 2(n - 1) data transfer operations over n links of bandwidth B to perform them on. So the time to perform one round of all-reduce can be estimated as ``` t = (S/B) * (2*(n-1)/n) ``` Consider networking with B = 1000Gbps = 125GB/second and n = 16, we have t = 0.42 seconds to perform one reduction. So we are network bound since the time to compute one step is longer than the time spent communicating between GPUs 0.29 seconds < 0.42 seconds. If we had chosen B = 3200Gbps = 400GB/second, t = 0.13 seconds and we enter the compute limited regime. Advanced algorithms will try to overlap the communication of gradients with the computation of later layers. Even if communication and computation perfectly overlapped every step, in the network bound case, there is at least 0.13 seconds of idling, which is almost 0.13 seconds / 0.42 seconds = 30% of the total step time! Very wasteful. ## NCCL Test NCCL (NVIDIA Collective Communications Library) is a library developed by NVIDIA to provide in collective communication primitives for GPU-GPU communication, which is used while training large models. If you've profiled your model training before, you've probably noticed the `ncclKernel_AllReduce` call. It's common to run the all reduce benchmark which is included as part of NVIDIA's [**nccl-test**](https://github.com/NVIDIA/nccl-tests/) to QA a GPU fabric. We ran this on a pair of 1600 Gbps machines, which showed the following: ``` out-of-place in-place # size count type redop root time algbw busbw #wrong time algbw busbw #wrong # (B) (elements) (us) (GB/s) (GB/s) (us) (GB/s) (GB/s) 8 2 float sum -1 29.73 0.00 0.00 0 29.65 0.00 0.00 0 16 4 float sum -1 28.43 0.00 0.00 0 28.02 0.00 0.00 0 32 8 float sum -1 28.60 0.00 0.00 0 28.48 0.00 0.00 0 64 16 float sum -1 29.02 0.00 0.00 0 29.01 0.00 0.00 0 128 32 float sum -1 29.14 0.00 0.01 0 29.13 0.00 0.01 0 256 64 float sum -1 29.47 0.01 0.02 0 29.56 0.01 0.02 0 512 128 float sum -1 31.29 0.02 0.03 0 30.83 0.02 0.03 0 1024 256 float sum -1 32.81 0.03 0.06 0 32.54 0.03 0.06 0 2048 512 float sum -1 35.57 0.06 0.11 0 35.71 0.06 0.11 0 4096 1024 float sum -1 37.27 0.11 0.21 0 36.39 0.11 0.21 0 8192 2048 float sum -1 38.52 0.21 0.40 0 38.15 0.21 0.40 0 16384 4096 float sum -1 42.18 0.39 0.73 0 41.70 0.39 0.74 0 32768 8192 float sum -1 43.50 0.75 1.41 0 42.53 0.77 1.44 0 65536 16384 float sum -1 48.09 1.36 2.56 0 45.70 1.43 2.69 0 131072 32768 float sum -1 56.49 2.32 4.35 0 55.72 2.35 4.41 0 262144 65536 float sum -1 103.7 2.53 4.74 0 93.15 2.81 5.28 0 524288 131072 float sum -1 97.47 5.38 10.09 0 97.48 5.38 10.08 0 1048576 262144 float sum -1 101.6 10.32 19.35 0 101.4 10.34 19.39 0 2097152 524288 float sum -1 115.6 18.15 34.03 0 110.6 18.97 35.57 0 4194304 1048576 float sum -1 148.7 28.21 52.88 0 149.5 28.06 52.62 0 8388608 2097152 float sum -1 223.3 37.57 70.44 0 221.7 37.84 70.96 0 16777216 4194304 float sum -1 303.1 55.35 103.77 0 303.8 55.23 103.55 0 33554432 8388608 float sum -1 511.7 65.58 122.96 0 510.9 65.68 123.15 0 67108864 16777216 float sum -1 953.1 70.41 132.03 0 945.6 70.97 133.07 0 134217728 33554432 float sum -1 1581.3 84.88 159.15 0 1578.5 85.03 159.43 0 268435456 67108864 float sum -1 2945.4 91.14 170.88 0 2950.6 90.98 170.58 0 536870912 134217728 float sum -1 5852.4 91.73 172.00 0 5821.6 92.22 172.91 0 1073741824 268435456 float sum -1 11567 92.83 174.05 0 11580 92.72 173.86 0 2147483648 536870912 float sum -1 22518 95.37 178.81 0 22437 95.71 179.46 0 # Out of bounds values : 0 OK # Avg bus bandwidth : 48.8808 ``` For training large models, we really only care the value of the busbw in the last row (which in this case is 178.81GB/s) since during training, gradients are bucketed and sent in large payloads such as the 2147483648 = 2GB message size in the last row. It's possible to achieve better efficiency through further tuning, with 95% of the theoretical bandwidth commonly achieved. Usually network tuning will just be handled by your cloud provider, but due diligence and silent faults with GPU fabrics say that you should health check your fabric periodically. ## Torch DDP Throughput Next, a quick DDP test helps ensure we can efficiently train models end-to-end. To iterate through datasets faster, Distributed Data Parallel (DDP) is commonly used to shard batches across data parallel workers by copying the model onto each GPU and communicating gradients between workers. To make this easier to test across clouds, we added this [**torch DDP benchmark to Skypilot**](https://github.com/skypilot-org/skypilot/tree/master/examples/torch_ddp_benchmark). After running this, we can get the following summary. ``` Benchmark: resnet101 with batch size 32 sec/iter ex/sec sec/iter ex/sec sec/iter ex/sec sec/iter ex/sec 1 GPUs -- no ddp: p50: 0.064s 501/s p75: 0.064s 499/s p90: 0.064s 497/s p95: 0.065s 491/s 1 GPUs -- 1M/1G: p50: 0.064s 502/s p75: 0.064s 502/s p90: 0.064s 502/s p95: 0.064s 501/s 2 GPUs -- 1M/2G: p50: 0.066s 486/s p75: 0.066s 486/s p90: 0.066s 484/s p95: 0.066s 482/s 4 GPUs -- 1M/4G: p50: 0.068s 468/s p75: 0.069s 464/s p90: 0.070s 457/s p95: 0.077s 417/s 8 GPUs -- 1M/8G: p50: 0.069s 465/s p75: 0.069s 464/s p90: 0.069s 463/s p95: 0.069s 463/s 16 GPUs -- 2M/8G: p50: 0.089s 359/s p75: 0.090s 356/s p90: 0.091s 350/s p95: 0.094s 340/s ``` Here we see that our throughput is not scaling ideally linearly, which is indicative of poor communications bandwidth. ## Conclusion The dream of any ML engineer is linear scaling. 2x the GPUs should mean half the training time. However, the only way to practically achieve this is with networking capabilities that can communicate gradients as faster than your GPU can compute them. Is your team or organization struggling with distributed training? I'd love to chat and see if Trainy can help! Reach out at [**roanak@trainy.ai**](mailto:roanak@trainy.ai). --- ## Blog: GPU Utilization is a Misleading Metric (https://www.trainy.ai/blog/gpu-utilization-misleading-metric) Published 2025-02-24 by Roanak Baviskar. The most common metric ML teams use to understand GPU usage is GPU Utilization, commonly accessed by running `nvidia-smi` from your terminal. Many integrated observability tools track GPU Utilization as their main performance metric as well. Somewhat to our surprise, it turns out that this is not always the best metric for understanding your GPU performance. In fact, you can get 100% GPU utilization by just reading/writing to memory while doing 0 computations! This article is about how we found that out, and what else we learned along the way. At [Trainy](https://www.trainy.ai/), we work on infrastructure for managing GPU clusters, so we spend a lot of time thinking about this stuff. Last year we were working with a foundation model company to scale out and improve the efficiency of their LLM training. We went through the basic steps that are mentioned in practically every guide about performance tuning for Pytorch — namely: - Saturate the GPU by changing dataloader defaults (`num_workers`, `batch_size`, `pin_memory`, prefetch factor, etc.) - Maximize tensor core usage by using mixed precision (fp16, bf16) - Use a fused optimizer from apex/deepspeed (e.g. FusedAdam, FusedAdamW, etc.) - Use instances/networking designed for training (H100SXM, A100SXM). Also, use newer instances when possible H100 > A100 > V100 These simple changes got us to 100% GPU utilization and significant power draw, which is great! To check whether more could be done, we calculated the MFUs of the training workload. > **A quick refresher: MFUs, or Model FLOPS (Floating point Operations Per Second) utilization, is one of the best metrics to understand GPU performance, as introduced in [Google's PaLM paper](https://arxiv.org/pdf/2204.02311). It is "the ratio of the observed throughput (tokens-per-second) relative to the theoretical maximum throughput of a system operating at peak FLOPs." In simpler terms, it tells you how many floating point operations per second your workload computes compared to your GPU's maximum capabilities. Its only real drawback is that MFU can be somewhat difficult to calculate compared to a metric like GPU utilization, since it is parameter and framework-dependent.** Unfortunately, the model training was only hitting ~20% MFUs. For reference, most LLM trainings today hit around [**35% - 45% MFUs**](https://github.com/mosaicml/llm-foundry/tree/main/scripts/train/benchmarking). So the question became: How could we only be using 20% of the theoretical max of our GPU's computational power, while also sitting at 100% GPU utilization? To answer this, we need a better understanding of what GPU utilization is actually tracking. ## What is GPU Utilization, really? GPU utilization is loosely defined on [**the Nvidia docs**](https://developer.nvidia.com/management-library-nvml) as "Current utilization rates are reported for both the compute resources of the GPU and the memory interface." This is wonderfully vague. A better definition can (surprisingly) be found on [**Datadog's NVML docs**](https://docs.datadoghq.com/integrations/nvml/#metrics), "Percent of time over the past sample period during which one or more kernels was executing on the GPU." To determine why this is misleading, we need a quick primer on how GPUs work. A GPU has [**cores and multiprocessing managers**](https://cvw.cac.cornell.edu/gpu-architecture/gpu-characteristics/kernel_sm). With Nvidia GPUs these multiprocessing managers are referred to as streaming multiprocessors (SMs), and on AMD hardware these are referred to as compute units (CU). Below is an illustration of the GH100 GPU, with 144 SMs. ![Illustration of H100 GPU, with 144 SMs](/blog/images/h100-sms-144.png) These multiprocessing managers can be thought of as foremen for a group of workers, in this case, cores. When you launch a CUDA kernel, the work is executed on CUDA cores by one or more SMs. As you can see below, a single SM on the GH100 chip has many CUDA cores. ![A single SM on an H100, with many CUDA cores](/blog/images/h100-single-sm.png) This means the metric, GPU Utilization, is only measuring whether a kernel is executing at a given time. It has no indication of whether your kernel is using all cores available, or parallelizing the workload to the GPU's maximum capability. **In the most extreme case, you can get 100% GPU utilization by just reading/writing to memory while doing 0 FLOPS.** Now we'd like to clarify: this is only misleading to someone without a background in systems (like many ML engineers). As mentioned [**here**](https://arthurchiao.art/blog/understanding-gpu-performance/#24-the-use-methodology), the definition of GPU utilization does make some sense under the [**"USE" methodology**](https://www.brendangregg.com/usemethod.html). But back to the problem at hand, this definition certainly explains the gap in GPU utilization % vs MFU % that we were seeing! There was certainly more performance left on the table, we just had to find it. ## Digging Deeper The next step to search for more performance was certainly to profile the model's training loop. We took a look at the training loop in Pytorch Profiler to get a better idea. As you can see below, the Softmax kernel was registering high GPU utilization, but low for a metric called SM efficiency. Now this was already sounding the alarm for us because naive softmax is a notorious bottleneck for LLMs, with many [**kernel fusions**](https://triton-lang.org/main/getting-started/tutorials/02-fused-softmax.html#motivations) such as [**FlashAttention**](https://github.com/Dao-AILab/flash-attention) coming out to address its memory-bound nature. Given this information, the SM efficiency statistic could be pointing out inefficiencies in our model's execution. ![Pytorch Profiler shows us low SM Efficiency during Softmax Kernel](/blog/images/profiler-sm-efficiency.webp) ## But What does SM Efficiency represent? SM efficiency (also known as SM activity) is a metric on Nvidia GPUs describing what % of SMs are active in a given time interval. As we previously mentioned, SMs can be thought of as foremen for a group of CUDA cores. An [**Nvidia H100 GPU**](https://developer.nvidia.com/blog/nvidia-hopper-architecture-in-depth/), for example, has 132 SMs, with 128 cores per SM, giving us a total of 16,896 cores. By measuring SM efficiency, we can determine whether our CUDA kernels are using our streaming multiprocessors. If we have a CUDA kernel that continuously runs for 10 seconds but only uses 1 SM, on an H100, this would register 100% utilization, but the SM efficiency would be 1 / 132 = 0.7%. Great, this is exactly what we were looking for! We can monitor SM efficiency layer by layer to determine what the low-hanging fruits are in terms of potential gain from optimization. ## Making Optimizations Now that we can easily identify which kernels aren't running hot on the GPU, we can work on optimizing these layers. Since this is a transformer stack, most gains are going to be made by fusing the layers in the transformer block definition. The figure below summarizes what we optimized. ![Kernel Fusions within Transformer Block](/blog/images/kernel-fusion-transformer.webp) By fusing, we mean instead of using a PyTorch native definition of a set of layers, we replace it with a GPU kernel, implemented in either CUDA or Triton, that combines all the layers into one kernel. The speedup results from less time per kernel reading/writing to GPU memory than time spent doing math in the case of certain layers (e.g. [**Softmax**](https://triton-lang.org/main/getting-started/tutorials/02-fused-softmax.html)). [**Flash Attention**](https://github.com/Dao-AILab/flash-attention) is an example of such a fused kernel. The other kernels that needed to be fused are the [**MLP**](https://github.com/Dao-AILab/flash-attention/blob/9a11f440d3a34f618b4ba814c825b109c6d7e8f5/flash_attn/ops/fused_dense.py#L531) and [**dropout layer norm residual add**](https://github.com/Dao-AILab/flash-attention/blob/9a11f440d3a34f618b4ba814c825b109c6d7e8f5/flash_attn/ops/fused_dense.py#L531) operations. Did we write these kernels ourselves? Nope. Most of these are already implemented in libraries like Flash Attention have implementations of layers as `nn.Modules` so that you don't have to worry about implementing a `torch.autograd.function` from scratch using the kernels. In addition, these implementations are often already hardware optimized so they use less memory in addition to just being faster. The biggest challenge is identifying where in your code you need to swap the appropriate layers. While `torch.compile` attempts to do this automagically, as of the time of writing this post, `torch.compile` [**doesn't play nicely with newer distributed strategies like FSDP**](https://dev-discuss.pytorch.org/t/torch-compile-fsdp-dec-8th/1718) and doesn't provide a lot of promised speedups in practice due to graph breaks. Hopefully, in the future, torch compilers can just do this for us, but for now, we are stuck manually adding fused implementations. For results, we achieved a 4x speedup in training time and 38% MFU for this customer, from the original MFU of 20%. Most of the optimizations we made came from these fused kernels, as well as finding the right "level" of model parallelism given their model size and the 3.2 Tbps InfiniBand that was available. ## Conclusion We'd highly recommend most AI teams to track SM Efficiency on their GPU cluster as well as GPU Utilization. It gives a much more representative picture of how much performance you're squeezing out of your GPUs, while GPU utilization can be a measure of whether or not the machine is idle. Of course, calculating MFUs would be great as well, but it isn't really a metric you can monitor all the time and layer-by-layer. Meanwhile, [**Nvidia DCGM**](https://docs.nvidia.com/datacenter/dcgm/latest/user-guide/feature-overview.html#profiling-metrics) (Data Center GPU Manager) ships SM activity by default. There are also more fine-grained metrics, such as SM occupancy (or Achieved Occupancy in Pytorch Profiler), that tell us how much work is being done by each SM. However, understanding these metrics is less straightforward than trying to get SM Efficiency as high as possible. If you're interested in learning more, I'd recommend taking a look at the [**Pytorch Profiler blog**](https://pytorch.org/blog/pytorch-profiler-1.9-released/#gpu-metric-on-timeline), [**DCGM docs**](https://docs.nvidia.com/datacenter/dcgm/latest/user-guide/feature-overview.html#profiling-metrics), [**Nsight's kernel profiling guide**](https://docs.nvidia.com/nsight-compute/ProfilingGuide/index.html), and [**Nsight docs**](https://docs.nvidia.com/gameworks/content/developertools/desktop/analysis/report/cudaexperiments/kernellevel/achievedoccupancy.htm). Thanks for giving this a read, and good luck squeezing every bit of performance out of your GPUs!