Uladzislau Sazanovich - JetBrains AI: Deep Dive

A deep dive into the architecture and implementation of JetBrains AI – the backbone of AI Assistant in JetBrains IDEs, Grazie, and other JetBrains products. The JetBrains AI platform covers a broad range of functionality, from coding to writing assistance. We will discuss both JetBrains’s own models and the caveats of integration with some well-known LLM providers. We will also shed light on inference techniques used in JetBrains AI products.

1. JetBrains AI Deep Dive Uladzislau Sazanovich uladzislau.sazanovich@jetbrains.com 1
2. Structure • JetBrains AI Overview • Code Completion Overview • Case Study: One-line completion 2
3. What are the characteristics of IDE features that come to mind? 3
4. Non-deterministic 4
5. Test generation 5
6. Commit generation 6
7. Built-in things 7
8. AI Chat 8
9. Code completion 9
10. Code Completion Overview 10
11. fun testGeneration() { val lm = LanguageModel() val prefix = "Once upon a time " 11
12. fun testGeneration() { val lm = LanguageModel() val prefix = "Once upon a time " Next word prediction, using prefix 12
13. fun testGeneration() { val lm = LanguageModel() val prefix = "Once upon a time " Language modeling! 13 Next word prediction, using prefix
14. fun testGeneration() { val lm = LanguageModel() val prefix = "Once upon a time " 14
15. fun testGeneration() { val lm = LanguageModel() val prefix = "Once upon a time " val expected = "is a golden retriever" val generation = lm.predict(prefix) assertEquals(generation, expected) } 15
16. fun testGeneration() { val lm = LanguageModel() val prefix = "Once upon a time " val expected = "is a golden retriever" val generation = lm.predict(prefix) assertEquals(generation, expected) } Doesn’t make any sense 16
17. fun testGeneration() { val lm = LanguageModel() val prefix = "My dog " val expected = "is a golden retriever" val generation = lm.predict(prefix) assertEquals(generation, expected) } 17
18. fun testGeneration() { val lm = LanguageModel() val prefix = "My dog " val expected = "is a golden retriever" val generation = lm.predict(prefix) assertEquals(generation, expected) } Prefix 18
19. fun testGeneration() { val lm = LanguageModel() val prefix = "My dog " val expected = "is a golden retriever" val generation = lm.predict(prefix) assertEquals(generation, expected) } Suffix Prefix 19
20. fun testGeneration() { val lm = LanguageModel() val prefix = "My dog " val expected = "is a golden retriever" val generation = lm.predict(prefix) assertEquals(generation, expected) } file: src/test/com/dogs/retriever/Generate.kt 20
21. fun testGeneration() { val lm = LanguageModel() val prefix = "My dog " val expected = "is a golden retriever" val generation = lm.predict(prefix) assertEquals(generation, expected) } file: src/test/com/dogs/retriever/Generate.kt Suffix Prefix One-line completion 21
23. ./ Example of LanguageModel usage fun runGeneration() { | 23
24. Multi-line completion ./ Example of LanguageModel usage fun runGeneration() { val lm = LanguageModel() val prefix = "My dog is " val generation = lm.predict(prefix) } 24
25. Multi-line completion ./ Example of LanguageModel usage fun runGeneration() { val lm = LanguageModel() val prefix = "My dog is " val generation = lm.predict(prefix) } Much harder task 25
27. Prompt 27 Type Declarations Local Variables Base Class Relevant Files MORE!
28. Prompt We want all of it 28 Type Declarations Local Variables Base Class Relevant Files MORE!
29. Prompt But we can’t 16k tokens Type Declarations Local Variables Base Class Relevant Files MORE! 29
30. Prompt 30
31. Prompt We use on-device ML to rank 31
32. Prompt We use on-device ML to rank As well a good old heuristics 32
33. IDE Multi-line context fun runGeneration() { val lm = LanguageModel() val prefix = "My dog is " val generation = lm.predict(prefix) } 33
34. IDE LanguageModel.kt class LanguageModel() { • fun predict(prefix: String) { ... } } Multi-line context fun runGeneration() { val lm = LanguageModel() val prefix = "My dog is " val generation = lm.predict(prefix) } 34
35. IDE LanguageModel.kt class LanguageModel() { • fun predict(prefix: String) { ... } } Multi-line context fun runGeneration() { val lm = LanguageModel() val prefix = "My dog is " val generation = lm.predict(prefix) } ProjectInfo file: Run.kt lang: Kotlin 1.9 libraries: ... 35
37. Cascade One-line ☁ Multi-line ☁ LLM ☁ 37
38. Cascade Local 🦙 One-line ☁ Multi-line ☁ LLM ☁ 38
39. Cascade Local 🦙 One-line ☁ Multi-line ☁ LLM ☁ 39
40. Case Study: One-line completion 40
41. fun testGeneration() { val lm = LanguageModel() val prefix = "My dog " val expected = "is a golden retriever" val generation = lm.predict(prefix) assertEquals(generation, expected) } Just one line 41
42. LLMs are great • Intelligent • Huge context • Easy to prototype 42
43. LLMs are great • Intelligent • Huge context • Easy to prototype • Slow and expensive 43
44. Training from Scratch • Flexibility 44
45. Training from Scratch • Flexibility • Full control of data 45
46. Training from Scratch • Flexibility • Full control of data • Any architecture 46
47. First try NVIDIA/ Megatron-LM starcoder 47
48. starcoder-15B First try 48
49. starcoder-15B starcoder-1B First try 49
50. starcoder-15B starcoder-1B First try 50 starcoder-160M
51. starcoder-15B starcoder-1B jetcoder-350M First try 51 starcoder-160M
52. val prefix = "My dog is golden " val expected = "retriever" val prefix = "My dog is golden " val exp = 1e6 MATCH MISS :( Line Accuracy 52
53. Line Accuracy, % 53
54. Let’s scale 54
55. Let’s scale 55 …a bit
57. 57 • Modern architecture • Developed by META • Supported in MegatronLM
58. 560ms 7 000 000 000 parameters 58
59. 560ms 590ms 350 000 000 parameters 7 000 000 000 parameters 59
60. size model has the same latency x20 60
61. • Added support for Starcoder architecture in DeepSpeed Optimise 61
62. • Added support for Starcoder architecture in DeepSpeed 590ms 326ms Optimise 62
63. ~0.5B parameters 63
64. ~0.5B parameters 6 layers wide-llama 64
65. ~0.5B parameters 24 layers tall-llama 65
66. class LanguageModel { fun predict(prefix: String) { ... } } fun testGeneration() { val lm = LanguageModel() val prefix = "My dog " val expected = "is a golden retriever" val generation = lm.predict(prefix) assertEquals(generation, expected) } 66
67. class LanguageModel { fun predict(prefix: String) { ... } } fun testGeneration() { val lm = LanguageModel() val prefix = "My dog " val expected = "is a golden retriever" val generation = lm.predict(prefix) assertEquals(generation, expected) } ⭐ 72 tokens 67
68. class LanguageModel { fun predict(prefix: String) { ... } } fun testGeneration() { val lm = LanguageModel() val prefix = "My dog " val expected = "is a golden retriever" val generation = lm.predict(prefix) assertEquals(generation, expected) } ⭐ 72 tokens 🦙 96 tokens 68
69. \n\n\t\t\t\t ⭐ 1 token 🦙 6 tokens 69
70. Overall • 🦙 0.5B • Tokeniser from starcoder • Tall and wide versions 70
71. Line Accuracy, % wide-llama-0.5B tall-llama-0.5B starcoder-160M starcoder-1B jetcoder-350M 71
72. Latency on A100 GPU, ms 72
73. We have a lot of 🦙🦙🦙! 73
74. Work in Progress 🔨🔨🔨 74
75. Thank you for listening! Uladzislau Sazanovich uladzislau.sazanovich@jetbrains.com 75