Search

We introduce MoMa, a novel modality-aware mixture-of-experts (MoE) architecture designed for pre-training mixed-modal, early-fusion language models. MoMa processes images and text in arbitrary sequences by dividing expert modules into modality-specific groups. These groups exclusively process designated tokens while employing learned routing within each group to maintain semantically informed adaptivity. Our empirical results reveal substantial pre-training efficiency gains through this modality-specific parameter allocation. Under a 1-trillion-token training budget, the MoMa 1.4B model, featuring 4 text experts and 4 image experts, achieves impressive FLOPs savings: 3.7x overall, with 2.6x for text and 5.2x for image processing compared to a compute-equivalent dense baseline, measured by pre-training loss. This outperforms the standard expert-choice MoE with 8 mixed-modal experts, which achieves 3x overall FLOPs savings (3x for text, 2.8x for image). Combining MoMa with mixture-of-depths (MoD) further improves pre-training FLOPs savings to 4.2x overall (text: 3.4x, image: 5.3x), although this combination hurts performance in causal inference due to increased sensitivity to router accuracy. These results demonstrate MoMa's potential to significantly advance the efficiency of mixed-modal, early-fusion language model pre-training, paving the way for more resource-efficient and capable multimodal AI systems., Comment: v2 -> update related work section v3 -> fix spelling

Scaling laws play an instrumental role in the sustainable improvement in model quality. Unfortunately, recommendation models to date do not exhibit such laws similar to those observed in the domain of large language models, due to the inefficiencies of their upscaling mechanisms. This limitation poses significant challenges in adapting these models to increasingly more complex real-world datasets. In this paper, we propose an effective network architecture based purely on stacked factorization machines, and a synergistic upscaling strategy, collectively dubbed Wukong, to establish a scaling law in the domain of recommendation. Wukong's unique design makes it possible to capture diverse, any-order of interactions simply through taller and wider layers. We conducted extensive evaluations on six public datasets, and our results demonstrate that Wukong consistently outperforms state-of-the-art models quality-wise. Further, we assessed Wukong's scalability on an internal, large-scale dataset. The results show that Wukong retains its superiority in quality over state-of-the-art models, while holding the scaling law across two orders of magnitude in model complexity, extending beyond 100 GFLOP/example, where prior arts fall short., Comment: 12 pages

We study a mismatch between the deep learning recommendation models' flat architecture, common distributed training paradigm and hierarchical data center topology. To address the associated inefficiencies, we propose Disaggregated Multi-Tower (DMT), a modeling technique that consists of (1) Semantic-preserving Tower Transform (SPTT), a novel training paradigm that decomposes the monolithic global embedding lookup process into disjoint towers to exploit data center locality; (2) Tower Module (TM), a synergistic dense component attached to each tower to reduce model complexity and communication volume through hierarchical feature interaction; and (3) Tower Partitioner (TP), a feature partitioner to systematically create towers with meaningful feature interactions and load balanced assignments to preserve model quality and training throughput via learned embeddings. We show that DMT can achieve up to 1.9x speedup compared to the state-of-the-art baselines without losing accuracy across multiple generations of hardware at large data center scales.

Two-level atoms coupled with single-mode cavity photons are predicted to exhibit a quantum phase transition when the coupling strength exceeds a critical value, entering a phase in which atomic polarization and photonic field are finite even at zero temperature and without external driving. However, this phenomenon, the superradiant phase transition (SRPT), is forbidden by a no-go theorem due to the existence of the diamagnetic term in the Hamiltonian. Here, we present spectroscopic evidence for a magnonic SRPT in ErFeO$_3$, where the role of the photonic mode (two-level atoms) in the photonic SRPT is played by an Fe$^{3+}$ magnon mode (Er$^{3+}$ spins). The absence of the diamagnetic term in the Fe$^{3+}$-Er$^{3+}$ exchange coupling ensures that the no-go theorem does not apply. Terahertz and gigahertz magnetospectroscopy experiments revealed the signatures of the SRPT -- a kink and a softening, respectively, of two spin-magnon hybridized modes at the critical point.

We present and benchmark a quantum computing approach to calculate the two-dimensional coherent spectrum (2DCS) of high-spin models. Our approach is based on simulating their real-time dynamics in the presence of several magnetic field pulses, which are spaced in time. We utilize the adaptive variational quantum dynamics simulation (AVQDS) algorithm for the study due to its compact circuits, which enables simulations over sufficiently long times to achieve the required resolution in frequency space. Specifically, we consider an antiferromagnetic quantum spin model that incorporates Dzyaloshinskii-Moriya interactions and single-ion anisotropy. The obtained 2DCS spectra exhibit distinct peaks at multiples of the magnon frequency, arising from transitions between different eigenstates of the unperturbed Hamiltonian. By comparing the one-dimensional coherent spectrum with 2DCS, we demonstrate that 2DCS provides a higher resolution of the energy spectrum. We further investigate how the quantum resources scale with the magnitude of the spin using two different binary encodings of the high-spin operators: the standard binary encoding and the Gray code. At low magnetic fields both encodings require comparable quantum resources, but at larger field strengths the Gray code is advantageous. Numerical simulations for spin models with increasing number of sites indicate a polynomial system-size scaling for quantum resources. Lastly, we compare the numerical 2DCS with experimental results on a rare-earth orthoferrite system. The observed strength of the magnonic high-harmonic generation signals in the 2DCS of the quantum high-spin model aligns well with the experimental data, showing significant improvement over the corresponding mean-field results.

We present a comprehensive theory of light-controlled multi-band superconductivity and apply it to predict distinctive signatures of light-driven superconducting (SC) states in terahertz multi-dimensional coherent spectroscopy (THz-MDCS) experiments. We first derive gauge-invariant Maxwell-Bloch equations for multi-band BCS superconductors. For this, we go beyond previously considered Anderson pseudo-spin precession models to include quantum transport effects. By calculating the THz-MDCS spectra measured experimentally, we then identify unique signatures of finite-momentum Cooper-pairing states that live longer than the laser pulse. These non-equilibrium SC states are characterized by long-lived canting of Anderson pseudo-spins. The pseudo-spin oscillators that describe the properties of these SC states are parametrically driven by both finite-momentum Cooper pairing and by time oscillations of the order parameter relative phase. We show that such strong parametric driving leads to drastic changes in the THz-MDCS spectral shape from the predictions of third-order nonlinear susceptibility calculations. These spectral changes strongly depend on the interband-to-intraband interaction ratio and on the collective modes of the light-driven state., Comment: 27 pages, 10 figures

The dynamics of Higgs mode in superconductors, manifested as coherent oscillations of the superconducting order parameter amplitude, provides vital insights into the nature of the superconducting gap structure and symmetry. Here we utilize two-dimensional terahertz coherent spectroscopy to investigate Higgs dynamics of a newly discovered infinite-layer nickelate superconductor. While we observe distinct nonlinear terahertz responses from the superconducting state, well-defined long-lived Higgs modes, as commonly observed in $s$-wave superconductors, are entirely absent in the nickelate film. Instead, we find the coherent nonlinear terahertz response is dominated by the quasiparticle excitations. These observations strongly indicate that the Higgs mode in infinite-layer nickelates is heavily damped by the quasiparticle excitations at arbitrarily low energies, which is a characteristic of $d$-wave pairing symmetry. Additionally, by examining the temperature dependence of the nonlinear terahertz response, we discover short-range superconducting fluctuations in the vicinity of $T_\mathrm{c}$. Our findings provide proof of a new $d$-wave system and establish a foundation for investigating the unconventional superconductivity in nickelates., Comment: 10 pages, 5 figures

The discovery of superconductivity in infinite-layer nickelates establishes a new category of unconventional superconductors that share structural and electronic similarities with cuprates. Despite exciting advances, such as the establishment of a cuprate-like phase diagram and the observation of charge order and short-range antiferromagnetic fluctuation, the key issues of superconducting pairing symmetry, gap amplitude, and superconducting fluctuation remain elusive. In this work, we utilize static and ultrafast terahertz spectroscopy to address these outstanding problems. We demonstrate that the equilibrium terahertz conductivity and nonequilibrium terahertz responses of an optimally Sr-doped nickelate film ($T_c$ = 17 K) are in line with the electrodynamics of $d$-wave superconductivity in the dirty limit. The gap-to-$T_c$ ratio 2$\Delta/k_\mathrm{B}T_\mathrm{c}$ is extracted to be 3.4, indicating the superconductivity falls in the weak-coupling regime. In addition, we observed significant superconducting fluctuation near $T_\mathrm{c}$, while it does not extend into the deep normal state as optimally hole-doped cuprates. Our result highlights a new $d$-wave system which closely resembles the electron-doped cuprates, expanding the family of unconventional superconductivity in oxides., Comment: 8 pages, 4 figures

The aim of this research is to reconstruct the 3D X-ray refractive index gradient maps by the proposed vector Radon transform and its inverse, assuming that the small-angle deviation condition is met. Theoretical analyses show that the X-ray beam can be modeled as a streamline with continuous change of direction in a row when measured in one grating period, which allows the extraction of the refraction angle signals. Experimental results show that all the 2D refraction signals of different directions can be acquired by a standard circular scanning procedure, which is typically used in the X-ray differential phase-contrast computed tomography. Furthermore, the 3D refractive index gradient maps that contain the directional density changes, can also be accurately reconstructed., Comment: accepted by IGTA2023

A recently emerging concept for quantum phase discovery is the controlled gapping of linear band crossings in topological semimetals. For example, achieving topological superconducting and charge-density-wave (CDW) gapping could introduce Majorana zero modes and axion electrodynamics, respectively. Light engineering of correlation gaps in topological materials provides a new avenue of achieving exotic topological phases inaccessible by conventional tuning methods such as doping and straining. Here we demonstrate a light control of correlation gaps and ultrafast phase switchings in a model CDW and polaron insulator (TaSe$_4$)$_2$I recently predicted to be an axion insulator. Our ultrafast terahertz photocurrent spectroscopy reveals a two-step, non-thermal melting of polarons and electronic CDW gap via studying the fluence dependence of a {\em longitudinal} circular photogalvanic current. The helicity-dependent photocurrent observed along the propagation of light reveals continuous ultrafast switchings from the polaronic state, to the CDW (axion) phase, and finally to a hidden Weyl phase as the pump fluence increases. Other distinguishing features corroborating with the light-induced switchings include: mode-selective coupling of coherent phonons to polaron and CDW modulation, and the emergence of a {\em non-thermal} chiral photocurrent above pump threshold of CDW-related phonons. The ultrafast chirality control of correlated topological states revealed here is important to realize axion electrodynamics and quantum computing., Comment: 9 pages, 4 figures

Generalized linear models (GLMs) are a widely utilized family of machine learning models in real-world applications. As data size increases, it is essential to perform efficient distributed training for these models. However, existing systems for distributed training have a high cost for communication and often use large batch sizes to balance computation and communication, which negatively affects convergence. Therefore, we argue for an efficient distributed GLM training system that strives to achieve linear scalability, while keeping batch size reasonably low. As a start, we propose P4SGD, a distributed heterogeneous training system that efficiently trains GLMs through model parallelism between distributed FPGAs and through forward-communication-backward pipeline parallelism within an FPGA. Moreover, we propose a light-weight, latency-centric in-switch aggregation protocol to minimize the latency of the AllReduce operation between distributed FPGAs, powered by a programmable switch. As such, to our knowledge, P4SGD is the first solution that achieves almost linear scalability between distributed accelerators through model parallelism. We implement P4SGD on eight Xilinx U280 FPGAs and a Tofino P4 switch. Our experiments show P4SGD converges up to 6.5X faster than the state-of-the-art GPU counterpar.

Sharding a large machine learning model across multiple devices to balance the costs is important in distributed training. This is challenging because partitioning is NP-hard, and estimating the costs accurately and efficiently is difficult. In this work, we explore a "pre-train, and search" paradigm for efficient sharding. The idea is to pre-train a universal and once-for-all neural network to predict the costs of all the possible shards, which serves as an efficient sharding simulator. Built upon this pre-trained cost model, we then perform an online search to identify the best sharding plans given any specific sharding task. We instantiate this idea in deep learning recommendation models (DLRMs) and propose NeuroShard for embedding table sharding. NeuroShard pre-trains neural cost models on augmented tables to cover various sharding scenarios. Then it identifies the best column-wise and table-wise sharding plans with beam search and greedy grid search, respectively. Experiments show that NeuroShard significantly and consistently outperforms the state-of-the-art on the benchmark sharding dataset, achieving up to 23.8% improvement. When deployed in an ultra-large production DLRM with multi-terabyte embedding tables, NeuroShard achieves 11.6% improvement in embedding costs over the state-of-the-art, which translates to 6.6% end-to-end training throughput improvement. To facilitate future research of the "pre-train, and search" paradigm in ML for Systems, we open-source our code at https://github.com/daochenzha/neuroshard, Comment: Accepted by MLSys 2023. Code available at https://github.com/daochenzha/neuroshard

Self-distillation relies on its own information to improve the generalization ability of the model and has a bright future. Existing self-distillation methods either require additional models, model modification, or batch size expansion for training, which increases the difficulty of use, memory consumption, and computational cost. This paper developed Self-discipline on multiple channels(SMC), which combines consistency regularization with self-distillation using the concept of multiple channels. Conceptually, SMC consists of two steps: 1) each channel data is simultaneously passed through the model to obtain its corresponding soft label, and 2) the soft label saved in the previous step is read together with the soft label obtained from the current channel data through the model to calculate the loss function. SMC uses consistent regularization and self-distillation to improve the generalization ability of the model and the robustness of the model to noisy labels. We named the SMC containing only two channels as SMC-2. Comparative experimental results on both datasets show that SMC-2 outperforms Label Smoothing Regularizaion and Self-distillation From The Last Mini-batch on all models, and outperforms the state-of-the-art Sharpness-Aware Minimization method on 83% of the models.Compatibility of SMC-2 and data augmentation experimental results show that using both SMC-2 and data augmentation improves the generalization ability of the model between 0.28% and 1.80% compared to using only data augmentation. Ultimately, the results of the label noise interference experiments show that SMC-2 curbs the tendency that the model's generalization ability decreases in the late training period due to the interference of label noise. The code is available at https://github.com/JiuTiannn/SMC-Self-discipline-on-multiple-channels.

It is widely acknowledged that large models have the potential to deliver superior performance across a broad range of domains. Despite the remarkable progress made in the field of machine learning systems research, which has enabled the development and exploration of large models, such abilities remain confined to a small group of advanced users and industry leaders, resulting in an implicit technical barrier for the wider community to access and leverage these technologies. In this paper, we introduce PyTorch Fully Sharded Data Parallel (FSDP) as an industry-grade solution for large model training. FSDP has been closely co-designed with several key PyTorch core components including Tensor implementation, dispatcher system, and CUDA memory caching allocator, to provide non-intrusive user experiences and high training efficiency. Additionally, FSDP natively incorporates a range of techniques and settings to optimize resource utilization across a variety of hardware configurations. The experimental results demonstrate that FSDP is capable of achieving comparable performance to Distributed Data Parallel while providing support for significantly larger models with near-linear scalability in terms of TFLOPS.

Laser-driven coherent phonons can act as modulated strain fields and modify the adiabatic ground state topology of quantum materials. Here we use time-dependent first-principles and effective model calculations to simulate the effect of the coherent phonon induced by strong terahertz electric field on electronic carriers in the topological insulator ZrTe$_5$. We show that a coherent $A_\text{1g}$ Raman mode modulation can effectively pump carriers across the band gap, even though the phonon energy is about an order of magnitude smaller than the equilibrium band gap. We reveal the microscopic mechanism of this effect which occurs via Landau-Zener-St\"uckelberg tunneling of Bloch electrons in a narrow region in the Brillouin zone center where the transient energy gap closes when the system switches from strong to weak topological insulator. The quantum dynamics simulation results are in excellent agreement with recent pump-probe experiments in ZrTe$_5$ at low temperature., Comment: 15 pages, 6 figures

We have developed a versatile near-field microscopy platform that can operate at high magnetic fields and below liquid-helium temperatures. We use this platform to demonstrate an extreme terahertz (THz) nanoscope operation and to obtain the first cryogenic magneto-THz time-domain nano-spectroscopy/imaging at temperatures as low as 1.8 K and magnetic fields of up to 5 T simultaneously. Our cryogenic magneto-THz scanning near-field optical microscopy, or cm-SNOM, instrument comprises three main equipment: i) a 5 T split pair magnetic cryostat with a custom made insert for mounting SNOM inside; ii) an atomic force microscope (AFM) unit that accepts ultrafast THz excitation and iii) a MHz repetition rate, femtosecond laser amplifier for high-field THz pulse generation and sensitive detection. We apply the cm-SNOM to obtain proof of principle measurements of superconducting and topological materials. The new capabilities demonstrated break grounds for studying quantum materials that requires extreme environment of cryogenic operation and applied magnetic fields simultaneously in nanometer space, femtosecond time, and terahertz energy scales.

Learning feature interactions is important to the model performance of online advertising services. As a result, extensive efforts have been devoted to designing effective architectures to learn feature interactions. However, we observe that the practical performance of those designs can vary from dataset to dataset, even when the order of interactions claimed to be captured is the same. That indicates different designs may have different advantages and the interactions captured by them have non-overlapping information. Motivated by this observation, we propose DHEN - a deep and hierarchical ensemble architecture that can leverage strengths of heterogeneous interaction modules and learn a hierarchy of the interactions under different orders. To overcome the challenge brought by DHEN's deeper and multi-layer structure in training, we propose a novel co-designed training system that can further improve the training efficiency of DHEN. Experiments of DHEN on large-scale dataset from CTR prediction tasks attained 0.27\% improvement on the Normalized Entropy (NE) of prediction and 1.2x better training throughput than state-of-the-art baseline, demonstrating their effectiveness in practice.

Classic models of cell size control consider cells divide while reaching a threshold, e.g. size, age, or size extension. The molecular basis of the threshold involves multiple layers of regulation as well as gene noises. In this work, we study cell cycle as first-passage problem with stochastic threshold and discover such stochasticity affects the inter-division statistics, which bewilders the criteria to distinguish the types of size control models. The analytic results show the autocorrelation in the threshold can drive a sizer model to the adder-like and even timer-like inter-division statistics, which is supported by simulations. Following the picture that the autocorrelation in the threshold can propagate to the inter-division statistics, we further show that the adder model can be driven to the timer-like one by positive autocorrelated threshold, and even to the sizer-like one when the threshold is negatively autocorrelated. This work highlights the importance to examine gene noise in size control., Comment: 9 pages, 9 figures

Image resolution has a significant effect on the accuracy and computational, storage, and bandwidth costs of computer vision model inference. These costs are exacerbated when scaling out models to large inference serving systems and make image resolution an attractive target for optimization. However, the choice of resolution inherently introduces additional tightly coupled choices, such as image crop size, image detail, and compute kernel implementation that impact computational, storage, and bandwidth costs. Further complicating this setting, the optimal choices from the perspective of these metrics are highly dependent on the dataset and problem scenario. We characterize this tradeoff space, quantitatively studying the accuracy and efficiency tradeoff via systematic and automated tuning of image resolution, image quality and convolutional neural network operators. With the insights from this study, we propose a dynamic resolution mechanism that removes the need to statically choose a resolution ahead of time.

Emergent superconductivity is strongly correlated with the symmetry of local atomic configuration in the parent compounds of iron-based superconductors. While chemical doping or hydrostatic pressure can change the local geometry, these conventional approaches do not provide a clear pathway in tuning the detailed atomic arrangement predictably, due to the parent compounds complicated structural deformation in the presence of the tetragonal-to-orthorhombic phase transition. Here, we demonstrate a systematic approach to manipulate the local structural configurations in BaFe2As2 epitaxial thin films by controlling two independent structural factors orthorhombicity (in-plane anisotropy) and tetragonality (out-of-plane/in-plane balance) from lattice parameters. We tune superconductivity without chemical doping utilizing both structural factors separately, controlling local tetrahedral coordination in designed thin film heterostructures with substrate clamping and bi-axial strain. We further show that this allows quantitative control of both the structural phase transition, associated magnetism, and superconductivity in the parent material BaFe2As2. This approach will advance the development of tunable thin film superconductors in reduced dimension.

ML workloads are becoming increasingly popular in the cloud. Good cloud training performance is contingent on efficient parameter exchange among VMs. We find that Collectives, the widely used distributed communication algorithms, cannot perform optimally out of the box due to the hierarchical topology of datacenter networks and multi-tenancy nature of the cloudenvironment.In this paper, we present Cloud Collectives , a prototype that accelerates collectives by reordering theranks of participating VMs such that the communication pattern dictated by the selected collectives operation best exploits the locality in the network.Collectives is non-intrusive, requires no code changes nor rebuild of an existing application, and runs without support from cloud providers. Our preliminary application of Cloud Collectives on allreduce operations in public clouds results in a speedup of up to 3.7x in multiple microbenchmarks and 1.3x in real-world workloads of distributed training of deep neural networks and gradient boosted decision trees using state-of-the-art frameworks.

Graph neural networks (GNNs), an emerging deep learning model class, can extract meaningful representations from highly expressive graph-structured data and are therefore gaining popularity for wider ranges of applications. However, current GNNs suffer from the poor performance of their sparse-dense matrix multiplication (SpMM) operator, even when using powerful GPUs. Our analysis shows that 95% of the inference time could be spent on SpMM when running popular GNN models on NVIDIA's advanced V100 GPU. Such SpMM performance bottleneck hinders GNNs' applicability to large-scale problems or the development of more sophisticated GNN models. To address this inference time bottleneck, we introduce ES-SpMM, a cache-first edge sampling mechanism and codesigned SpMM kernel. ES-SpMM uses edge sampling to downsize the graph to fit into GPU's shared memory. It thus reduces the computation cost and improves SpMM's cache locality. To evaluate ES-SpMM's performance, we integrated it with a popular GNN framework, DGL, and tested it using representative GNN models and datasets. Our results show that ES-SpMM outperforms the highly optimized cuSPARSE SpMM kernel by up to 4.35x with no accuracy loss and by 45.3x with less than a 1% accuracy loss.

Deep learning recommendation models (DLRMs) are used across many business-critical services at Facebook and are the single largest AI application in terms of infrastructure demand in its data-centers. In this paper we discuss the SW/HW co-designed solution for high-performance distributed training of large-scale DLRMs. We introduce a high-performance scalable software stack based on PyTorch and pair it with the new evolution of Zion platform, namely ZionEX. We demonstrate the capability to train very large DLRMs with up to 12 Trillion parameters and show that we can attain 40X speedup in terms of time to solution over previous systems. We achieve this by (i) designing the ZionEX platform with dedicated scale-out network, provisioned with high bandwidth, optimal topology and efficient transport (ii) implementing an optimized PyTorch-based training stack supporting both model and data parallelism (iii) developing sharding algorithms capable of hierarchical partitioning of the embedding tables along row, column dimensions and load balancing them across multiple workers; (iv) adding high-performance core operators while retaining flexibility to support optimizers with fully deterministic updates (v) leveraging reduced precision communications, multi-level memory hierarchy (HBM+DDR+SSD) and pipelining. Furthermore, we develop and briefly comment on distributed data ingestion and other supporting services that are required for the robust and efficient end-to-end training in production environments.