Apple Silicon용 MLX 사용하기

안녕하세요 저는 Awni입니다 오늘 MLX를 소개하게 되어 기쁩니다 MLX는 Apple Silicon용으로 제작된 오픈 소스 배열 프레임워크입니다 유연한 도구로 기본적인 수치 계산부터 Apple 기기에서 최첨단 머신 러닝 모델 실행까지 폭 넓게 활용할 수 있습니다 대규모 언어 모델을 사용하여 텍스트를 생성하려면 최신 모델을 사용하여 이미지와 오디오, 비디오를 생성하세요 MLX가 도와드립니다 MLX로 학습 및 미세 조정도 가능하며 Mac에서 직접 머신 러닝 모델을 사용자화할 수도 있죠 먼저 MLX와 그 유용성을 자세히 설명해보죠 Python에서 MLX의 기본 사용법도 다룰텐데 여기엔 설치 및 기본 배열 연산이 포함됩니다 그리고 몇 가지 주요 기능을 소개하겠습니다 MLX를 다른 프레임워크와 차별화하는 요소죠 그리고 MLX가 제공하는 몇 가지 툴을 알아보겠습니다 Apple Silicon에서 머신 러닝 워크로드의 빠른 처리를 돕죠

마지막 섹션에서는 MLX Swift와 그 사용법을 소개하겠습니다 MLX 소개부터 시작하겠습니다 MLX는 완전히 새롭게 구축되어 Apple Silicon에서의 빠른 실행과 CPU에서의 실행도 가능하며 GPU에서 가속 실행도 가능하죠 다양한 응용 프로그램에 MLX를 사용할 수 있습니다 소규모 수치 연산부터 대규모 머신 러닝까지요 사용하기 쉽고 유연하도록 설계되었으며 속도와 효율성을 저하시키지 않습니다

MLX는 NumPy와 유사한 핵심 API를 가집니다 드롭인 대체로 사용해서 대부분의 수치 계산을 가속화합니다 MLX는 자동 미분 상위 수준 라이브러리 등 머신 러닝에 필요한 모든 툴도 갖추고 있죠 이와 같은 상위 수준 API는 PyTorch 및 JAX와 유사합니다 언급된 프레임워크 중 하나를 사용했다면 MLX를 쉽게 익히고 시작할 수 있을 겁니다 기기에서 MLX를 사용해 더 발전된 머신 러닝을 수행할 수 있죠 예를 들어, LM Studio에서는 MLX를 사용하는데요 LM Studio는 유명한 텍스트 생성 앱으로 Mac에서 바로 대규모 언어 모델로 텍스트를 생성할 수 있죠

MLX 위에 구축된 MLX LM 패키지를 사용해 수백억개 매개변수 규모의 언어 모델로 텍스트를 생성하고 미세 조정할 수 있습니다 더 자세한 내용을 알아보려면 'MLX를 사용하여 Apple Silicon에서 대규모 언어 모델 탐색하기'를 시청하세요 MLX는 빠른 프로토타입 제작이 가능한 Python API를 제공하죠 Swift에도 API가 있는데 신경망을 구축하고 최적화하기 위한 모든 상위 레벨 패키지가 포함됩니다 MLX는 C++와 C 언어용 API도 제공합니다 어떤 언어로든 MLX를 사용해 Apple Silicon에서 최신 머신 러닝 모델을 실행할 수 있습니다 여기에는 Mac, iPhone, iPad Vision Pro가 포함되죠 MLX 소프트웨어는 제약이 적은 MIT 오픈소스 라이선스를 따릅니다 핵심 소프트웨어는 GitHub에서 확인할 수 있는데 Python과 Swift API로 빌드한 예제 및 패키지와 함께 제공됩니다 MLX는 Hugging Face에 활발한 모델 제작자 커뮤니티도 보유하고 있죠 최신 모델 중 다수는 이미 MLX 커뮤니티의 Hugging Face 조직에 있고 새 모델이 매일 업로드됩니다 PyPi에서 설치하면 Python에서 MLX를 가장 쉽게 사용할 수 있죠 터미널에서 pip3 install mlx 한 줄만 실행하면 됩니다 MLX는 사용이 간편하죠 배열에서 연산 작업을 실행하려면 Python 파일을 열고 MLX를 가져오기만 하면 됩니다 그러면 배열을 만들고 기본적인 연산을 수행할 수 있죠 예를 들어 두 개의 정수 배열을 추가합니다

배열에 대한 정보를 쉽게 검사할 수도 있죠 예를 들어 모양과 데이터 유형 등이요 앞서 말했듯 MLX Python API는 NumPy와 비슷합니다 보통 작업 이름과 서명이 같고 동일한 방식으로 작동하죠 NumPy 또는 유사한 프레임워크를 사용했다면 MLX를 쉽게 익히고 시작할 수 있을 겁니다 이제 MLX가 무엇인지, 어디에 유용한지와 기본 사항을 알았으니 주요 기능에 대해 알아보겠습니다 여기에는 통합 메모리와 지연 평가, 함수 변환 신경망을 구축 및 최적화하기 위한 상위 레벨 패키지가 있습니다 MLX는 Apple Silicon 기능을 최대한 활용하도록 설계되었는데 통합 메모리에 특화된 새로운 프로그래밍 모델도 그중 하나죠

일반적으로 머신 러닝에 사용되는 대부분의 시스템에는 독립형 GPU와 별도 메모리가 탑재됩니다 그러나 Apple Silicon은 통합 메모리 아키텍처를 사용합니다 CPU와 GPU가 동일한 물리적 메모리를 공유한다는 뜻이죠

통합된 메모리로 작업하려면 MLX는 익숙한 기존 프레임워크와 다른 방식을 사용합니다 기존 프레임워크에서는 연산은 데이터를 따릅니다 배열이 CPU 메모리에 있으면 연산은 CPU에서 이뤄지죠 배열이 GPU 메모리에 있으면 연산은 GPU에서 이뤄집니다

MLX에서 배열은 통합 메모리에 할당됩니다 지원되는 기기에서 사용하기 위해 복사할 필요가 없습니다

대신 기기에서 연산을 실행하려면 연산에 사용할 기기를 지정합니다 예를 들어 a와 b를 GPU에서 더하고 a와 b를 CPU에서 곱합니다 작업은 병렬로 실행될 수 있으며 의존성이 존재하면 MLX가 자동으로 관리합니다 MLX의 또 다른 주요 기능은 지연 소비입니다 MLX의 효율을 극대화하기 위해 특히 소비가 큰 경우 MLX는 지연 실행 엔진을 사용하죠

덧셈과 같은 두 배열 간 연산이 호출되면 실제 계산은 일어나지 않습니다 대신 여기 보이는 것처럼 계산 그래프가 구축됩니다

배열 C는 아직 계산되지 않았습니다 실제로 사용해야 하는 경우에만 계산됩니다 예를 들어, C를 인쇄하거나 MLX에서 C를 Python 리스트로 변환하면 강제로 계산이 실행되죠 mx.eval을 사용해 그래프 평가를 명시적으로 강제할 수 있습니다 지연 계산에는 여러 이점이 있습니다 빌드 과정을 분리하고 계산 그래프를 실행해 MLX는 결과를 계산하기 전에 그래프에서 변환과 최적화가 가능합니다 또한, 지연 계산을 사용하면 사용한 만큼만 자원이 소모되죠 함수 변환은 MLX의 또 다른 주요 특징입니다 빠른 배열 프레임워크를 훨씬 더 강력한 툴로 발전시켜 머신 러닝 모델을 학습시키고 성능을 개선합니다

함수 변환은 함수를 입력받아 새 함수로 반환하는 변환입니다

MLX에는 여러 가지 함수 변환이 있습니다 보통 다음 두 가지 카테고리 중 하나에 속하죠 자동 미분을 위한 변환은 계산 그래프를 최적화하는 작업입니다 예를 들어 함수 변환을 사용해 자동으로 MLX에서 함수의 기울기를 계산할 수 있죠

입력값의 사인을 계산하는 기본 함수가 있다고 해보죠 이 함수의 기울기를 구하려면 mx.grad 함수 변환을 사용하면 됩니다 mx.grad의 출력은 새로운 함수로 배열에서 호출하면 미분값을 나타냅니다 함수 변환은 임의 조합이 가능합니다 단순히 mx.grad를 mx.grad의 출력에 사용해서 2차 미분이 가능합니다

결과는 또 다른 함수이며 배열에 적용했을 때 부호 함수의 2차 미분값을 나타내죠

MLX에는 두 가지 상위 레벨 패키지도 있어 신경망을 쉽게 구축하고 학습시킬 수 있습니다

첫 번째는 신경망 구축에 사용되는 모듈식 라이브러리 mlx.nn입니다 두 번째는 일반적 최적화 알고리즘 라이브러리 mlx.optimizers입니다 두 패키지는 독립적으로 사용할 수 있지만 서로 완벽하게 통합됩니다

mlx.nn 패키지는 신경망을 구축하는 데 필요한 모든 기능을 갖추고 있습니다 nn.module은 레이어와 컨테이너가 모두 상속받는 기본 클래스입니다 매개변수 액세스, 매개변수 로드 매개변수 저장 등 유용한 방법을 제공합니다

nn 라이브러리에는 기본 제공 표준 레이어 세트도 함께 제공되죠 예를 들어 nn.Linear같은 것들이요 nn.Module에서 상속받아 나만의 레이어를 만들 수도 있습니다

일반적으로 사용되는 손실 함수 및 초기화 방법은 nn.losses 및 nn.init 하위 패키지에도 포함됩니다

mlx.nn을 사용해 간단한 다층 신경망을 구축하는 방법을 알아보겠습니다

우선 nn.Module에서 상속받는 맞춤형 클래스를 만듭니다 간단한 단일 은닉층 신경망을 사용하겠습니다 모듈의 초기화 메서드 내부에 선형 레이어를 생성합니다 그리고 call 함수를 구현합니다 주어진 입력에 대한 모듈의 출력을 계산하는 함수죠 call 함수는 첫 번째 선형 레이어를 호출하고 Relu 활성화 함수를 적용한 다음 두 번째 선형 레이어를 호출합니다 MLX는 Apple Silicon의 통합 메모리 아키텍처에 맞게 설계 및 최적화되어 쉬운 시작을 도우며 상위 수준 신경망 패키지 또한 PyTorch같은 일반적인 머신 러닝 프레임워크와 유사하도록 설계되었습니다

MLX 모델 구현과 PyTorch의 동일한 구현을 비교해보죠 둘은 거의 동일합니다 출력을 계산하는 함수에만 작은 차이점 두 가지가 있죠 이전에 PyTorch에서 모델을 빌드한 적 있다면 MLX로 전환하는 것은 매우 간단할 것입니다 이제 MLX의 핵심 기능 대부분을 살펴보았으니 머신 러닝 워크로드를 가장 빠르고 효과적으로 처리하는 방법에 대해 이야기해 보죠 먼저 함수를 컴파일하여 속도를 높이는 방법을 보여드리겠습니다 그리고 mx.fast 하위 패키지를 소개할텐데 일반적인 머신 러닝 연산을 바로 사용할 수 있는 고속 구현체입니다 그리고 맞춤형 Metal 커널 추가용 API도 다룹니다 그 다음에는 양자화를 사용해 메모리 사용량을 줄이고 모델 실행 속도를 높이는 법도 알려드리죠 마지막으로 MLX를 사용해 여러 머신에 계산 작업을 분산하는 방법도 보여드립니다 MLX의 현실적인 계산 대부분은 배열에 대해 다수의 연산을 수행하는 함수로 구성됩니다 이와 같은 함수를 더 빠르게 만드는 간단한 방법은 mx.compile 함수 변환을 사용하는 거죠 몇 가지 요소별 연산을 수행하는 함수가 있다고 가정해 보겠습니다 여기 보이는 GELU 활성화 함수처럼요 이 함수의 계산 그래프에는 여러 개의 노드가 있습니다 각 노드는 내부적으로 별도의 GPU 커널 실행에 대응합니다

그래프를 단일 융합 커널로 컴파일하는 데는 모든 개별 커널이 사용됩니다 이런 방식으로 메모리 대역폭과 그래프 실행 오버헤드를 절약하고 훨씬 더 효율적인 계산이 가능합니다

mx.compile을 사용하는 것은 mx.compile 함수 변환을 함수에 데코레이터로 적용하는 것만큼 쉽습니다 컴파일은 종종 잘 작동하지만 머신 러닝의 기본 연산이 포함된 더 복잡한 연산의 경우에는 mx.fast 하위 패키지를 사용하면 더 빨리 실행할 수 있을 것입니다 예를 들어, 변압기 모델의 핵심 구성 요소 중 다수는 mx.fast에서 연산을 수행합니다 여기에는 위치 인코딩과 정규화 계층 스케일 점곱 어텐션이 포함됩니다 mx.fast의 연산은 더 전문화되었지만 최대한 빠른 학습과 추론을 위해 고도로 조정되었습니다 또한 유연한 구성이 가능해서 주어진 계산의 다양한 변형을 가속화할 수 있습니다 예를 들어 스케일 점곱 주의 연산은 입력으로 선택적 마스크를 사용할 수 있죠 마스크는 가산 마스크나 부울 마스크 또는 마스크 유형을 나타내는 문자열입니다 mx.fast의 연산 중 하나인 RMS 정규화를 자세히 살펴보죠 최신 트랜스포머 기반 대규모 언어 모델 대부분은 RMS norm을 쓰죠 MLX 연산으로 기본 구현만 해도 계산 그래프가 매우 커집니다

대신 전체 구현을 단일 호출인 mx.fast.rms_norm으로 대체할 수 있습니다 코드는 더 간단하고 계산 그래프는 단일 노드만을 가지며 계산 자체도 훨씬 더 빨라질 것입니다

MLX에는 맞춤형 Metal 커널을 추가하기 위한 API가 있어 더 맞춤형으로 구현해 성능이 향상될 수 있죠 아직 mx.fast에는 포함되어 있지 않습니다 맞춤형 커널을 작성하면 나머지는 모두 MLX가 처리합니다 적시 컴파일 및 실행을 포함해서요 이 커널은 Metal로 작성됐는데 GPU에서 실행되도록 Apple이 만든 언어이자 API입니다 Metal 코드로 된 소스 문자열과 입력 및 출력 정보를 넘겨 커널을 빌드합니다

그리드 크기와 출력의 형태 및 유형을 지정해서 커널을 호출합니다 MLX는 커널 호출을 다른 연산과 동일한 방식으로 처리합니다 MLX는 계산 그래프에 노드를 생성하며 지연 평가됩니다

머신 러닝 워크로드를 빠르게 만들 툴킷의 다른 툴은 양자화입니다 대규모 모델은 많은 메모리와 높은 메모리 대역폭으로 속도를 내죠 대부분의 학습에서 같거나 거의 동일한 수준의 퀄리티를 얻으려면 추론보다 훨씬 높은 정밀도가 필요합니다 정밀도를 낮춰 메모리에 큰 모델을 저장하고 빠르게 실행할 수 있죠

모델이 32비트 크기의 부동소수점 정밀도로 표현되면 첫 단계로 bfloat16 또는 float16을 사용해서 메모리 필요량을 절반으로 줄일 수 있습니다 16비트가 너무 많으면 MLX에는 내장 루틴이 있으니 양자화 배열을 더 작게 만들고 그것으로 연산을 실행할 수 있습니다 예를 들어 요소당 4비트를 사용해 양자화해서 메모리 필요량을 더 줄일 수 있는 거죠

행렬을 양자화하기 위해 mx.quantize를 사용합니다 사용할 요소별 비트 수와 그룹 크기를 지정합니다 그룹 크기는 공통된 스케일과 바이어스 값을 공유하는 양자화된 행렬의 요소 수를 결정합니다 비트가 작을수록 그룹 크기가 더 클수록 결과가 더 작고 더 빠릅니다 MLX에는 비트와 그룹 크기에 대한 여러 옵션이 있어 최대한의 유연성을 제공합니다

mx.quantized_matmul을 이용해 양자화된 행렬과 양자화되지 않은 입력 벡터 및 행렬 값을 곱할 수 있습니다

mx.dequantize를 사용하여 원본 입력의 근사값을 복구하세요

MLX.mn의 편리한 유틸리티로 단일 명령 모듈 양자화도 가능하죠 당신의 모델이 임베딩 레이어와 뒤를 잇는 다수의 선형 레이어 스택으로 구성되었다고 한다면 nn.quantize를 사용하여 모델을 양자화할 수 있습니다 양자화 명령은 또한 선택적 콜백을 사용해 어떤 레이어를 양자화하고 주어진 레이어에 어떤 정밀도를 사용할지 더 세밀하게 제어할 수 있습니다

대규모 언어 모델로 텍스트를 생성할 때 양자화는 메모리 사용을 극적으로 줄이고 초당 생성되는 토큰을 늘립니다 한 대의 머신만으론 충분하지 않기도 합니다 예를 들어, 대규모 모델로 텍스트를 생성하고 싶을 수 있으나 단일 머신의 메모리에 맞지 않죠 혹은 대규모 데이터 셋에서 모델을 미세 조정하거나 평가할 수 있는데 둘 다 병렬화하기 쉬워서 여러 대의 머신을 사용하면 훨씬 더 빨라질 수 있죠

MLX는 기본적으로 여러 대의 머신에서 임의 계산을 수행할 수 있는 기능을 제공합니다 머신은 이더넷이나 Thunderbolt를 통해 연결할 수 있습니다

mx.distributed 하위 패키지를 사용해 여러 머신에 계산 작업을 분산하세요 MX 분산은 주로 일련의 통신 연산입니다

예를 들어, all_Sum은 모든 머신에 입력 배열을 추가합니다 all_Sum의 출력은 합산된 입력입니다 각 기계마다 동일하죠 다수의 머신에서 배열을 합산하는 방법을 자세히 살펴보겠습니다 mx.distributed.init을 사용하여 분산 백엔드를 초기화합니다 이 단계는 선택 사항이지만 통신 그룹에 접근해야 할 때 호출하는 함수입니다 통신 그룹에는 유용한 정보가 첨부되어 있는데 총 프로세스 수와 현재 프로세스 인덱스 등입니다

그런 다음 각 프로세스에 단일 값을 갖는 배열을 만들고 mx.distributed.all_sum을 호출해 모든 프로세스의 배열을 합칩니다

다수의 머신에서 MLX 프로그램 실행을 돕는 런처가 있습니다 4대의 머신에서 프로그램을 실행하려면 4개의 호스트 IP 주소와 함께 mlx.launch를 사용합니다 지금까지 Python에서 MLX를 사용하는 방법을 봤는데요 대부분의 경우 편리하고 유연한 Python을 선호할 것입니다 아니면 Swift를 선호할 수도 있죠 그래서 MLX는 Swift기반의 모든 기능을 갖춘 API를 제공합니다 Metal을 기반으로 구축되었으며 macOS, iOS, iPadOS visionOS 등에서도 잘 적용됩니다 Swift에서 MLX를 시작하는 것은 MLX를 Xcode 프로젝트 패키지로 추가하는 것만큼 쉽습니다 프로젝트를 클릭한 다음 패키지 종속성 탭에서 더하기 기호를 클릭하세요 그런 다음 MLX Swift GitHub 저장소의 URL을 입력하고 패키지 추가 버튼을 클릭하세요 MLX Swift로 빌드를 시작하는 데 필요한 건 이게 전부입니다 Python과 Swift 간의 간편한 전환을 위해 두 언어 간의 API가 의도적으로 유사하게 설계되었죠 이전에 본 Python 코드 조각을 MLX Swift 대응 제품과 나란히 놓고 비교해보죠 배열을 만들고, 연산을 실행하며 배열에 대한 메타데이터 검사는 Swift에서 Python과 거의 동일하죠 MLX를 Python에서 사용하며 살펴본 모든 핵심 기능과 MLX 가속화 섹션에서 논의한 최적화는 MLX Swift에도 동일하게 적용됩니다 MLX의 많은 주요 기능의 광범위한 개요를 살펴보았습니다 프레임워크에 대해 더 알아보려면 MLX 웹사이트를 확인하세요 웹사이트에서 문서, 예제 등의 링크를 확인할 수 있습니다 Python과 Swift API 모두 예제 저장소가 있는데 일반적인 머신 러닝 사용 사례 예시가 많이 포함되어 있죠 언어 모델 학습 및 생성 이미지 생성, 음성 인식 등이 포함됩니다 예시를 통해 MLX에 대해 자세히 알아보고 MLX로 프로젝트를 구축하는 출발점으로 삼으세요 시청해 주셔서 감사합니다 MLX로 만들 당신의 다음 작품을 기대하겠습니다