Alert button
Picture for Siddhartha Rao Kamalakara

Siddhartha Rao Kamalakara

Alert button

Exploring Low Rank Training of Deep Neural Networks

Add code
Bookmark button
Alert button
Sep 27, 2022
Siddhartha Rao Kamalakara, Acyr Locatelli, Bharat Venkitesh, Jimmy Ba, Yarin Gal, Aidan N. Gomez

Figure 1 for Exploring Low Rank Training of Deep Neural Networks
Figure 2 for Exploring Low Rank Training of Deep Neural Networks
Figure 3 for Exploring Low Rank Training of Deep Neural Networks
Figure 4 for Exploring Low Rank Training of Deep Neural Networks

Training deep neural networks in low rank, i.e. with factorised layers, is of particular interest to the community: it offers efficiency over unfactorised training in terms of both memory consumption and training time. Prior work has focused on low rank approximations of pre-trained networks and training in low rank space with additional objectives, offering various ad hoc explanations for chosen practice. We analyse techniques that work well in practice, and through extensive ablations on models such as GPT2 we provide evidence falsifying common beliefs in the field, hinting in the process at exciting research opportunities that still need answering.

Viaarxiv icon
Github IconAccess Paper or Ask Questions
Twitter IconFacebook IconLinkedin IconWhatsapp IconMessenger Icon

Scalable Training of Language Models using JAX pjit and TPUv4

Add code
Bookmark button
Alert button
Apr 13, 2022
Joanna Yoo, Kuba Perlin, Siddhartha Rao Kamalakara, João G. M. Araújo

Figure 1 for Scalable Training of Language Models using JAX pjit and TPUv4
Figure 2 for Scalable Training of Language Models using JAX pjit and TPUv4
Figure 3 for Scalable Training of Language Models using JAX pjit and TPUv4
Figure 4 for Scalable Training of Language Models using JAX pjit and TPUv4

Modern large language models require distributed training strategies due to their size. The challenges of efficiently and robustly training them are met with rapid developments on both software and hardware frontiers. In this technical report, we explore challenges and design decisions associated with developing a scalable training framework, and present a quantitative analysis of efficiency improvements coming from adopting new software and hardware solutions.

* 5 pages, 4 figures 
Viaarxiv icon
Github IconAccess Paper or Ask Questions
Twitter IconFacebook IconLinkedin IconWhatsapp IconMessenger Icon