마이크로소프트의 Agent Lightning이 2026년 11월 3일 저정밀도 추론을 쓰는 강화학습 예제와 메모리 개선을 담은 새 버전을 공개했다. 에이전트 학습 개발자는 낮은 정밀도로 답을 생성할 때 학습을 안정화하는 보정 방식을 실험하고, 여러 전문가를 나눠 쓰는 모델에서는 학습 경로를 재생할 때 주 메모리 부담을 줄일 수 있다.
1. 저정밀도 추론에서 보상 학습 안정화 시도
이번 릴리스는 Qwen3.5-9B 수학 학습 예제에 답을 생성한 모델과 학습 모델의 확률 차이로 생기는 평균적인 학습 방향의 치우침을 빼는 보정인 스코어 센트링을 추가했다. 낮은 정밀도로 답을 만드는 모델과 학습 중인 모델은 다음 토큰을 고를 확률이 달라질 수 있다. 스코어 센트링은 이 차이 때문에 학습의 갱신 방향이 치우치는 문제를 겨냥한 보정 방식이다.
예제는 생성 위치마다 토큰을 뽑은 당시의 상위 128개 후보 확률을 보관하고, 현재 학습 모델의 확률과 대조해 보정값을 계산한다. 기록하지 않은 나머지 후보의 확률은 현재 모델에 비례한다고 근사한다.
회사 발표 기준 스코어 센트링은 후반 학습 보상을 약 0.6으로 유지했다. 응답 묶음의 보상을 비교해 학습하는 기존 GRPO를 보정 없이 쓴 기준선, 비교를 위한 기준 결과는 약 45번째 단계에 보상이 0으로 떨어진 뒤 유지됐다.
이 결과는 B200 한 장에서 Qwen/Qwen3.5-9B를 쓴 수학 예제의 학습 곡선, 학습에 따른 성능 변화다. 답 생성에는 INT8 가중치·활성화와 INT4 KV 캐시를 쓰고, 학습 연산은 BF16으로 진행했다. 보상은 math_verify가 최종 답을 판정해 정답이면 1, 아니면 0을 주며, 잘린 생성도 0점이다.
2. 실행 예제는 추론 정밀도와 학습 정밀도를 나눈다
릴리스에는 모델 준비와 정밀도 검증 도구, 학습 가중치를 INT8 추론 모델로 동기화하는 경로도 포함됐다. 양자화, 수치 표현의 비트 수를 줄이는 변환의 대상은 선형 계층, 입력을 가중합으로 바꾸는 신경망 층이며, 임베딩, 토큰을 수치 벡터로 바꾼 표현과 언어 모델 출력 계층은 BF16으로 남는다. 설정은 다음과 같이 추론 부분에 둔다.
actor_rollout_ref:
rollout:
dtype: bfloat16
quantization: int8
quantization_config_file: ${oc.env:SC_INT8_CONFIG,outputs/precision_validation/Qwen3.5-9B-W8A8-INT8/config.json}
engine_kwargs:
vllm:
linear_backend: triton
attention_backend: TRITON_ATTN
kv_cache_dtype: int4_per_token_headdtype은 주변 계산 정밀도를 정하고, 명시된 양자화·캐시 설정이 INT8과 INT4 경로를 선택한다. 예제의 전체 학습 설정은 프롬프트 32개에 응답 4개씩, 200회 업데이트이며, 빠른 점검은 프롬프트 2개와 응답 4개씩 3단계로 실행한다.
3. MoE 라우팅 재생은 임시 배열을 순차 처리
MoE는 여러 전문가 중 일부를 골라 토큰을 처리하는 모델 구조다. R3 라우팅 재생은 답 생성 때 선택한 전문가 정보를 학습 과정으로 전달해 같은 경로를 재생한다. 이전에는 묶음의 라우팅 배열을 모두 해독한 뒤 최종 텐서에 복사했지만, 이번 수정은 한 번에 하나의 배열만 해독하고 복사 후 해제한다.
개선 제안자가 공개한 로컬 CPU 측정은 합성 생성 자료 64개를 처리한 결과다. 아래 표처럼 라우팅 정보를 풀어 놓는 임시 공간과 프로세스의 최대 상주 메모리(RSS)가 모두 줄었다.
표: R3 라우팅 재생의 호스트 메모리
조건: 로컬 CPU 벤치마크, 합성 rollout, 모델이 생성한 응답 묶음 64개
단위: MB
| 항목 | 변경 전 | 변경 후 |
|---|---|---|
| 임시 디코딩 할당 피크 | 126.7 | 4.7 |
| 프로세스 피크 RSS | 660.3 | 543.4 |
| 최종 텐서 | 128 | 128 |
절감 폭은 임시 디코딩 버퍼에서 가장 크고, 최종 학습 텐서는 변경 전후 모두 128 MB다. 결과 텐서의 형식과 패딩·절단 동작을 유지하면서 중간 처리 공간을 줄인 개선이다.
4. 절감의 대가와 적용 범위
저정밀도 추론은 메모리·연산 효율을 노리지만, 스코어 센트링은 후보 확률을 기록해 전달하고 학습 모델의 어휘 전체 확률을 계산하는 일을 추가한다. 기록되지 않은 분포를 근사하는 방식이어서 보정값에는 오차 가능성이 남는다. 이 예제의 GRPO와 스코어 센트링은 그룹 표준편차 방식의 보상 정규화, 보상값을 그룹 내 변동 폭으로 조정하는 방식을 사용하며, 논문에 제시된 평균만 쓰는 설정 및 경사하강법, 기울기를 따라 손실을 낮추는 학습법과 동일한 실험 재현은 아니다.
릴리스에는 보상 소수값 보존과 정규화 정밀도 개선도 포함됐다. 보상을 BF16으로 먼저 바꿔 1.0과 1.001이 같은 값으로 처리되던 문제를 고쳤고, 정규화 분모의 오버플로, 표현 범위를 넘은 수치 오류와 토큰 수 반올림도 방지했다. 이는 강화학습 업데이트의 입력값이 낮은 정밀도 변환이나 큰 수 계산에서 사라지는 상황을 줄이는 수정이다.
이번 버전은 저정밀도 추론과 MoE 라우팅 재생의 병목을 각각 겨냥하며, 메모리 절감은 최종 텐서가 아닌 중간 버퍼 처리에서 나온다.