AI Infrastructure at Periodic

At Periodic, our infrastructure enables us to efficiently train specialized models that Pareto-dominate frontier models, including GPT-6 Astra and Claude Fable 5.1, on our X-ray diffraction evaluations. Starting from open weights, our final training run required only a peak of 1,300 H200 GPUs across the midtraining and reinforcement learning (RL) phases. These models now analyze experiments in our high-throughput labs, enabling us to search for better superconductors and magnets.

We make this possible through improvements in training throughput, inference generation speed, custom sandboxing, GPU memory efficiency, and scientific tool execution. Additionally, we share capacity with scientific modeling and simulation workloads in order to sustain 95%+ cluster utilization. Our goal is to make each GPU-hour go further, make our researchers maximally productive, and shorten the time between an idea to experimental results.

We’ve improved our efficiency against standard solutions across training, inference, and sandboxing. The bar plots highlight one metric for each, more details about the what and how can be found in the sections below. Together, these improvements let us test more ideas and deploy better models sooner.

We built on top of open-source frameworks including Megatron, SGLang, Miles, and Ray. From there, we optimize and heavily modify parts to handle scientific RL workloads, where RL rollouts can spend hours reasoning and running tools while training steps only take several minutes. To do so, we substantially improved open-source libraries for compute-efficient training and faster inference. In particular, our stack delivers 4.1x the training throughput of Megatron on the same GPUs, and our contributions to SGLang made inference 2.5x faster for our use case. We run training and inference asynchronously on separate GPU allocations, so both can progress as rollouts finish, and optimize how they process long, variable-length traces.

Example of a typical reasoning trace lasting over an hour on our X-ray diffraction analysis task, showing measured tool-call times and estimated inference times for Periodic Infra and OSS Infra (SGLang + SOTA Sandbox Provider). Tool-call times were measured by replaying a saved rollout; this tests real conditions, including possible differences in hardware.

We also make scientific tool execution and deployment far more efficient. Our sandboxes, called Periodic sandbox or pbox for short, isolate resource-intensive scientific tools to avoid destabilizing inference and training engines while using spare CPU capacity on GPU nodes. The RL job transfers files directly to and from its pbox sandboxes within the cluster, achieving 4.5x lower data transfer latency and 3.3x throughput of a state-of-the-art sandbox provider. For rapid experimentation, we reduced the time to convert Megatron checkpoints for inference from 30 minutes to one minute.

Periodic’s AI infrastructure has benefited from the years of development from the open-source community and we contribute our improvements back to Megatron-LM, SGLang, and Miles to support the open-source AI community.

Inference

Our inference stack is optimized for long-running RL. We adopt the PipelineRL approach: we run generation and training run concurrently and a single rollout may span multiple weight versions. Each RL rollout is a multi-turn conversation interleaved with tool calls and LLM generation. We use SGLang’s “in-place” mode to preserve the KV cache during weight synchronization and reuse it in subsequent generations. With sticky routing and if no eviction occurs, each turn can reuse the KV cache for all preceding tokens.

Overall architecture of our inference stack.

We cover the adopted features and specific inference improvements in the following sections.

WideEP Tensor parallelism (TP) with latent attention (MLA) is not ideal for high-throughput workloads. Because MLA compresses the KV cache into a single latent head, tensor-parallel ranks duplicate the KV cache.

WideEP uses data parallelism for attention and expert parallelism for mixture of experts (MoE), increasing effective KV-cache capacity by up to N times across N ranks. The tradeoff is latency: TP distributes attention computation across devices, while WideEP computes all attention heads on each DP rank. TP therefore provides better interactivity, but WideEP supports larger batches and reaches a higher throughput ceiling when KV-cache capacity is the bottleneck.

Router Replay

We use router replay to reduce the mismatch between rollout and training. After integrating SGLang router replay + WideEP into our RL stack, we observed a significant slowdown. Profiling pointed to longer CPU-side Gloo broadcasts. By merging profiles across ranks, we found the real issue: one rank was slowing down every DP rank. At long context lengths, returning routing data results in a large CPU overhead. Because DP-attention ranks synchronize during decoding, a slowdown on one rank propagates across the entire group.

In 32-way WideEP, any of the 32 ranks finishing a generation can trigger this synchronization. If the average output length is 200 tokens, each rank has an approximate 1/200 probability of finishing at a given step. The probability that at least one rank finishes is:

1 − (199/200)³² ≈ 15%

This means the slowdown can occur roughly 1.5 times every ten decoding steps!

The server originally returned routing decisions for the full sequence after every turn. We optimized this by implementing Delta router replay, which caches prior routing decisions on the client and requests only the newly generated suffix. The optimization was upstreamed to SGLang PR #24851.

Detailed explanation of the Delta router replay mechanism.

PD Disaggregation

When using WideEP, if any rank within a DP attention group is running prefill, the whole forward pass will not use CUDA graphs, which means decode on the other ranks will be significantly slowed down. We attempted to use prefill delayer in SGLang to delay prefill steps and avoid frequent prefill interruptions. However, when the number of ranks is high, such as 32, the chance of prefill is so high that prefill interruptions can still happen frequently. Instead, we adopt prefill-decode (PD) disaggregation that reduces rollout duration by 2x, making the rollouts less stale and reducing the required number of RL environments. We’ve upstreamed several PRs to streamline PD + RL.

The one drawback of PD disaggregation is KV-cache capacity, which is reduced due to redundancy across the P/D nodes. To mitigate this, we use decode-side radix cache to reduce the KV-cache transfer data between the prefill and decode nodes.

With PD disaggregation and delta router replay, we made decoding 2.5x faster, from 10 tok/s/request to 25 tok/s/request for a 1T model at the same batch size.

Fast Resharding

For trillion-parameter models, we reduced checkpoint resharding from Megatron to Hugging Face format from 30 minutes to 1 minute by distributing conversion across Ray workers and eliminating redundant full-tensor copies.

Weight resharding is embarrassingly parallel, so we distribute conversion tasks across Ray actors. For MoE weights, we added a direct path: instead of reconstructing the full destination tensor and copying loaded chunks into it, our converter reads file-local expert chunks and converts them directly into Hugging Face expert tensors. It was upstreamed to Miles PR #1371.

Fast weight resharding from Megatron to HuggingFace.

Training

We optimize our training stack for faster iterations and more experiments on the same GPU budget. The stack supports both large training runs and smaller RL trainers, with efficient handling of long, variable-length scientific traces.

Our stack achieves 4.1x the training throughput of our Megatron baseline¹ on long, variable-length inputs using the same GPUs. For the trillion-parameter Kimi K2.6, we reach around 350 TFLOP/s per GPU on Hopper at 256K context and support full-parameter training at 64K context on 64 H200s.

Scaling Up Efficiently: Larger Models and Longer Contexts

Long-context attention distributes work across tokens and heads. MoE distributes expert weights and routes tokens to the GPUs that own them. A GPU layout that suits one can be a poor fit for the other.

We use node-local overlapped TP and zigzag CP within and across nodes for attention, with PP distributing layers across stages. Within each stage, we fold TP and CP into EP for MoE computation. We use HybridEP and DeepEP for expert communication, fusing token permutation into dispatch. We tune the combined layout for each context length and cluster size.

The three columns show complementary optimizations: sequence packing, attention and MoE parallelism, and memory reduction. The 64-H200 trainer excludes GPUs used for rollout generation.

Scientific RL traces are typically long and vary widely in length. Padding wastes compute, while distributing tokens equally can still leave DP, PP, and CP workloads unbalanced.

Our THD sequence-packing implementation removes most padding while preserving each trace’s attention boundaries and loss masks. To balance work across data-parallel workers and pipeline microbatches, our cost-aware scheduler sorts traces by estimated cost, then greedily assigns each to the least-loaded microbatch with available capacity. This longest-processing-time-first approach reduces waits between DP workers and pipeline stalls from uneven microbatch workloads.

We partition each packed row once for CP, rather than each trace separately. Attention, loss computation, and routing replay share the same layout and global sample boundaries, reducing per-trace reshuffling while preserving attention isolation.

Traces are packed, balanced across workers and microbatches, then partitioned for CP. Zigzag CP pairs early and late chunks to balance causal-attention work.

We also demonstrated zero train–rollout log-probability difference in synchronous RL setup. In production, we use asynchronous RL, which is faster in practice and performs well.

Scaling Down Easily: Smaller RL Units, Higher Research Efficiency

Reducing GPU memory use lets us train on fewer GPUs and run with less recomputation. We shard the model with 3D parallelism (TP, PP, and EP) and distribute optimizer states across GPUs, then reduce the remaining memory footprint with selective recomputation, fusion, and FP8 primary weights. We retain activations that are expensive to recompute and recompute memory-heavy intermediates at low compute cost. Fused linear cross-entropy avoids materializing the full token-by-vocabulary logits tensor while preserving FP32 accumulation and token log-probabilities, and FP8 primary weights eliminate the redundant BF16 working-weight copy.

The remaining optimizer state is large but only needed during updates. We keep FP32 master weights and Adam states on CPU and stream them through a bounded GPU workspace. Transfers are pipelined across chunks, and Adam still runs on GPU. FP8 working weights stay resident for forward and backward; chunked checkpoint staging controls save-time peaks.

Sharding, chunked state offload, and selective recomputation reduce per-GPU memory. Compared with full recomputation, selective recomputation retains more activations to avoid repeating expensive operations.

At 64 GPUs and 64K context, these optimizations reduce the peak memory requirement from 320 GB to 132 GB. This enables full-parameter Kimi K2.6 training on 64 H200s without significant loss of per-GPU throughput, using half the GPUs of the smallest comparable state-of-the-art setup we know of. We collaborated with NVIDIA Devtech team to bring chunked optimizer offloading to Megatron-LM, this capability is now merged into Megatron’s dev branch.

Sandboxing

Scientific RL depends on executing model-written code efficiently and safely alongside expensive GPU workloads. Our first RL loop ran training, inference, and model-written code together, with no isolation. That worked until the model wrote code that requested 80 GB of RAM and took down the whole job. The open-source sandboxes we tried didn’t fit our cluster. Hosted providers offered the functionality we needed for scientific tasks, but became expensive at our scale.

So we built pbox, our own sandbox solution.

We framed pbox around one idea: run sandboxes on the nodes already allocated to the RL job by Slurm, instead of launching a separate sandbox service. Sandboxes use a separate CPU set from training and inference, and their lifetime follows the job’s allocation

This pbox design has several innovations that are different from industry-standard solutions:

• GPU-sandbox colocation. Training and inference leave substantial CPU capacity available on our GPU nodes. Pbox puts that capacity to work running CPU-intensive tools during rollouts.

• Slurm-based scheduling: Sandboxes live and die with the RL job. So scheduling is solved – no need to implement queueing, priority, and preemption in the sandboxing infra, no need to implement extreme auto scaling for RL.

• Local communication. The RL job communicates directly with pbox agents on its allocated nodes. In our test with 100 sandboxes, pbox completed 1 MiB upload-and-download round trips 4.5x faster at p50 and delivered 3.3x the aggregate throughput of the hosted provider we tested. Traffic stays within our cluster, avoiding transfers to and from an external provider.

• No stranded compute: When our RL job is not running, the Slurm job naturally releases the CPUs back to our compute pool, and it’s picked up by other jobs which don’t require sandboxes. In comparison, self-hosting a sandbox infra usually creates a stranded compute pool dedicated for sandboxing.

Top: comparing pbox with a state-of-the-art sandbox provider on startup time, 1 MiB upload-and-download round-trip latency (p50), and aggregate data throughput, using 100 sandboxes. Pbox traffic stays within our cluster network, whereas traffic to the provider leaves the cluster and returns. Thanks to the colocated sandbox architecture, pbox has lower latency and higher throughput. Bars show the average of 4 runs; error bars show the minimum and maximum. Bottom: speedups for sampled code executions in sampled code executions. Speedup is provider execution time divided by pbox execution time; values above 1 mean pbox is faster. The median per-execution speedup is 2.5x, and pbox is faster in 86.9% of sampled executions.

Deep integration of sandboxing with Slurm and Kubernetes

Pbox runs sandboxes alongside our RL job, on the same compute nodes. It is designed for Slurm-on-Kubernetes environments such as SUNK, so sandbox capacity follows the node resources Slurm has already allocated to the job.

At startup, the RL job divides CPUs on each allocated node into two separate sets: one for the RL workload and one for its sandboxes. It configures the pbox DaemonSet on that same node to use the sandbox CPU set. The job maintains a connection to pbox throughout execution. When the job exits or is preempted, the connection closes and Pbox cleans up its sandboxes.

The RL job then uses pbox to create gVisor sandboxes, execute commands, and transfer files. Sandboxes share some cached base images, while OverlayFS gives each sandbox its own writable layer.

As a result, pbox completely follows slurm’s scheduling, including preemption, priorities, and gang scheduling alongside RL jobs, even under heavy RL autoscaling. This also allows us to use many CPUs from our RL GPU nodes for sandboxing. In practice, we found using 32 CPUs per GPU node is usually sufficient.

Illustration of pbox’s architecture.

Looking Forward

The training, inference, and sandboxing improvements we described in this blog post have greatly accelerated our internal LLM work. We have no shortage of ideas for what to do next: multi-agent RL, stable lower-precision post-training quantization, distributed KV cache to tackle long rollout eviction, and scaling up sandboxes to thousands of images. A more open-ended problem in our roadmap is automatic calibration of computational resources across our diverse scientific computing setup that includes large databases, judge LLMs, atomistic simulations, and CPU-heavy scientific computations.

If you’re passionate about building state-of-the-art AI infrastructure for the physical sciences, explore our open positions.