西岸逐星：微软亚洲研究院半年实习感悟

Chaofan

2024-01-18

| Views: | Total Words: 4.8k | Reading Time: 4 mins.

前言

这篇文章改自我的本科阶段暑研实习报告，作为我在 MSRA 上海系统组这半年（实际算上远程，接近一年）实习体验分享。

面试

随着新年的钟声响起，我在微软亚洲研究院的半年多实习也接近了尾声。回想这段实习经历，精彩而又伴随着劳累。我总共给微软贡献了约有15000行代码，并且有一篇在投论文，这次实习我在学术和视野方面都收获多多。

在刚进入微软的时候，每个人都可以领到一个 “明日之星（Star of Tomorrow）” 的衣服，这是微软对实习项目的一个代号。再结合亚研院的位置，便有了这篇感悟的题目“西岸逐星”。这半年里，在我的mentor韩震华、杨玉庆的引导下，我逐渐从一个只有着代码经验的“打手”变成一个逐渐可以自己提出想法，研究问题要点的“研究者”。不过要逐到“星“，想必还有很长的路要走。

与亚研院的初识在他们来致远学院给的一个Talk之上。我印象很深刻，他们那时候分享的工作是他们在OSDI’22（系统领域旗舰会议）中稿的工作 SparTA，这篇工作用巧妙地将稀疏信息在张量的计算图上传播，使得稀疏深度学习推理得到了很大的加速。当时我刚好也在做相关的工作（计算图优化），听到这个工作顿生一种拍手叫绝之感，于是便萌生了暑研前往微软系统组实习的想法。于是在Talk结束后我主动上前和玉庆聊天，并且分享了我自己关于 SparTA 一些想法、改进点，之后也是成功地加到了他的联系方式。

不久以后，他们组告诉我说有一个Hiring Counter，于是我便参加了面试。面试并不考察技术性的问题，而是让我分享实验室实践做过的工作，我便把在APEX实验室和天奇学长合作时关于在Relax上做自动求导的工作介绍了一下。面试我的宁馨就是那篇 SparTA 的一作。最终我也是通过了面试，踏上了我的“逐星之旅“。

通勤与入职

不过首先，这个“旅途”从物理意义上就不算近——交大距离西岸有一小时半的地铁车程。因此我每次上班时都需要早早起来，到公司打卡的时候往往已经是上午九点多十点。后来我给这往返约3小时的挤地铁时光安排了一个活动：刷Paper。当然，地铁上不适合精读，因此我一般用手机刷一些学术Blog或者Arxiv的摘要，看到有趣的想法记下来，等回学校的时候再细读。一来二去，这成了我现在也一直保持的习惯，在外面乘车的时候也会抽空看看最新的学术消息。有一件很有趣的事就是我的通勤线路和迪士尼有一部分是重合的，因此我经常抬头能看见戴着米奇耳朵的游客。

进入微软的第一天就让我感受到浓浓的“大厂风“，实习生需要自己办理一系列手续（包括文书工作、微软内部账号注册、资源申请等等），这些微软都有一套很完整的软硬件支持，也有HR小姐姐给我们新人答疑解惑，但是有时候也让人感觉过于繁琐（比如很多公司的内部资源需要连内网或者通过VPN才能访问）。这都是大厂的特质，“庞大、厚重”，从正面看是“资源丰富、完整、成体系”，但是从另一面看也可以是“包袱重、手续繁琐”。微软这边的办公环境可以说是非常不错，不仅有一周五天各不相同的饮品，印着微软LOGO的马克杯，睡眠舱，升降桌和工学椅，更重要的就是他灵活的考勤制度。作为外企，它是一个注重Work Life Balance的地方，这样松弛的环境带来的反而是效率的极大提高。以及作为外企，它有着浓厚的西方文化气息，但又不失本土化。

研究经历

说完了大致情况，接下来就该讲讲这半年的实习历程了。刚进来时我还没确定现在的研究课题，而是帮他们之前的一个项目做一些性能优化，比如优化一个矩阵乘法的算子。究这样过了两三周，震华在会上说，他准备做一个大语言模型相关的项目，想法是将KV Cache在不同机器上高效地Migration。我们内部把这个项目命名为LLMOS（意思是LLM的操作系统），这个代号也沿用到了最后。为了做好这个Migration，我们首先需要大量的运行数据，包括OPT模型和Llama模型的每个token的延时。这个任务自然落到了我的肩上，于是我开始密集地测量，并且不断整理大量的数据。可惜的是，在一个月多的努力后，最后我们逐渐放弃了这个想法，因为数据表明我们能做的并不多。我在微软的第一段科研历程也就这样没有了后续。

不过我和震华一直试图在之前的主题（也就是LLMOS）上谋求新的点。我们逐渐将问题从General的LLM加速转向Agent的加速。在有一次我洗澡的时候，我突然想到，在Langchain的示例中有一些DAG形状的LLM任务，这种任务不同LLM请求之间实际上是存在“通信”的，也就是一个请求的输出会作为另一个请求的输入。因此我在组里面说出了我的想法：能不能在上一个请求不断生成token的时候，下一个请求同时进行填充，来达到优化通信的效果。

虽然这个想法后面测试出来的加速比不大，但是这给我们的研究方向打开了一个新的空间。我们逐渐将研究重心转向了“DAG形状的LLM任务“，并且尝试性地提出了一些新概念，比如”自然语言编程“、”LLM编译器“等。最后在大概九月中旬的时候，我们确定了论文的雏形：以一个Placeholder为核心，建立DAG，然后做优化。这时离会议的DDL还有不到三个月，于是我们紧锣密鼓地开始了工作的推进。初步的分工是我来完成整个代码的实现，然后震华来把Paper的基本框架先搭起来。这段时间大概是我最”充实“的一段时间了，因为工作内容完全就是在生产代码、和同事讨论接口，非常地具有成就感。

在十月多的时候，因为一些设计上的缺陷，我又对整个代码仓库进行了一次重构，几乎是把整个项目重写了一遍。同时，我还测量了一些关键数据，发现在common prefix很长的情况下，现有的方法L2 Cache的命中率特别低，I/O量特别大，因此我们这里可以实现一个kernel来进一步推进优化效果。不过很遗憾，我并不是写CUDA的专家；但是组里有一个同事很擅长写这个，于是我们找到他开始了合作。

在重构完之后，时间也基本来到了11月，这时我们意识到：Paper的实验需要开始做了。在有点赶的时间里，我开始了主动加班。为了节省时间，我申请了直接在学校工作，节省了3小时的车程，然后就是每天在包图从早上待到晚上。最后在会议的截止日期前，我和震华通宵到最后一刻，终于把所有实验补齐，并交上了我们的答卷。在Submitted出现的那一刻，我长舒一口气，此时窗外清晨的太阳早已倾泻进来。我和震华道了别，去学校的麦当劳吃了顿早餐，然后直接瘫在床上，开始补觉。这可以说是我实习最累的一刻，也是最放松的一刻了。

总结

之后就是围绕着这个项目的小修小补。总体来讲，我这次实习的故事线还是比较清晰，经历了一个完整项目的开发周期，并且也成功产出了论文，体验了一次完整的学术生产（从读别人Paper，到自己孵化Idea，再到不断修改最后成型）的过程。遗憾的是由于我忙于写代码，在Paper撰写方面得到的锻炼还不够多。以及还有几点从他们身上学到的东西、感想：

多读论文，多想，不要怕想法失败；
多和别人讨论，有时别人会看到你看不到的一些“盲点“；
多做实验，用数据说话，“Get hands dirty“。
Dogfooding。用自己的工作形成一个生态（可以是Follow-up，也可以是联动），扩大影响力。

以上就是我在微软的半年实习报告。微软亚研院系统组给我留下了非常好的印象，同事都十分有想法并且有行动力，这也是他们这几年Paper大丰收的原因吧。此外，微软的一些企业文化，内部召开的一些读书会、聚餐、来自海外名校教授的Talk也丰富了我的实习生活，让我感受到了微软的视野所在。
在这次大模型的风潮中，微软可谓是大赢家，不仅和OpenAI深入合作，还把自己的Copilot和Bing搞得风风火火。我能在这个时间点来到这里实习，也要感谢ACM班的良好声誉和丰富资源给到我这个机会，让我体验了和别人不一样的一个学期。最后再感谢一下微软特供的冬日养生茶，支持我最后阶段赶完这篇Paper。