언어에 관계없는 샤딩

컬렉션을 사용해 정리하기 내 환경설정을 기준으로 콘텐츠를 저장하고 분류하세요.

장기적인 목표는 Shardy를 모든 MLIR 방언과 호환되는 완전히 독립형 구성요소로 만드는 것입니다. 현재 Shardy는 StableHLO에 직접 종속되지만 Shardy를 더 유연하게 만들기 위해 다양한 추상화와 인터페이스를 통해 이를 제거하기 위해 노력하고 있습니다.

샤딩 규칙

샤딩 규칙은 작업을 통해 전파하는 방법을 인코딩합니다. 이제 Shardy는 StableHLO에 종속되므로 각 stablehlo 작업에 대한 샤딩 규칙을 정의합니다. 또한 Shardy는 방언 소유자가 작업에서 자체 작업의 샤딩 규칙을 정의하는 데 사용할 수 있는 ShardingRuleOpInterface를 제공합니다. 작업이 이 인터페이스를 구현하는 한 Shardy는 이를 통해 전파할 수 있습니다.

def ShardingRuleOpInterface : OpInterface<"ShardingRuleOpInterface"> {
  let methods = [
    InterfaceMethod<
      /*desc=*/[{
        Returns the sharding rule of the op.
      }],
      /*retType=*/"mlir::sdy::OpShardingRuleAttr",
      /*methodName=*/"getShardingRule"
    >,
  ];
}

데이터 흐름 작업

지역 기반 연산과 같은 일부 연산에는 모든 피연산자와 결과 간의 측정기준 간 대응 관계만 설명하는 샤딩 규칙만으로는 충분하지 않은 다른 접근 방식이 필요합니다. 이 경우 Shardy는 방언 소유자가 작업을 통해 샤딩의 전파를 설명할 수 있도록 ShardableDataFlowOpInterface를 정의합니다. 이 인터페이스는 소유자를 통해 각 데이터 흐름 가장자리의 소스와 타겟을 가져오고 가장자리 소유자의 샤딩을 가져오고 설정하는 메서드를 제공합니다.

def ShardableDataFlowOpInterface :
    OpInterface<"ShardableDataFlowOpInterface"> {
  (get|set)BlockArgumentEdgeOwnerShardings;
  (get|set)OpResultEdgeOwnerShardings;
  getBlockArgumentEdgeOwners;
  getOpResultEdgeOwners;
  getEdgeSources;
  // ...
}

데이터 흐름 작업을 처리하는 방법에 관한 대략적인 개요는 데이터 흐름 작업을 참고하세요.

아직 구현되지 않은 인터페이스

향후 Shardy를 더 유연하고 방언에 관계없는 언어로 만들기 위해 더 많은 인터페이스와 트레잇이 추가될 예정입니다. 아래에서 확인하세요.

상수 분할

MLIR의 대부분의 텐서 프로그램에는 이 값이 필요한 모든 연산자에 의해 재사용되는 상수 인스턴스가 하나 있습니다. 이는 필요한 상수가 동일한 경우에 적합합니다. 그러나 프로그램의 최적 샤딩을 위해 상수의 각 사용에 자체 샤딩이 있고 다른 연산에서 상수를 사용하는 방식의 영향을 받지 않도록 허용하고자 합니다.

예를 들어 아래 그림에서 add가 샤딩된 경우 계산의 여러 부분에 있는 divide 및 subtract가 샤딩되는 방식에 영향을 주지 않아야 합니다.

이를 거짓 종속 항목이라고 합니다. 상수는 비용이 적으므로 동일한 상수를 사용하는 연산 간에 실제 종속 항목이 없습니다. 따라서 사용자는 상수 (및 상수와 유사한) 연산의 샤딩을 결정할 수 있습니다. 그러면 이 상수를 사용할 때마다 고정된 하위 계산의 자체 사본에 개별적으로 전파될 수 있는 다른 샤딩이 있을 수 있습니다.

이를 위해 Shardy 사용자는 다음을 정의해야 합니다. - your_dialect.constant -> sdy.constant 전달 - sdy::ConstantLike 트레잇(예: iota) - add 및 multiply와 같은 요소별 연산을 위한 mlir::Elementwise 트레잇 - slice/broadcast와 같은 연산을 위한 sdy::ConstantFoldable 이러한 연산은 모든 피연산자/결과가 상수인 경우 컴파일 시 기술적으로 계산될 수 있습니다.

작업 우선순위

GSPMD에서는 요소별 연산이 먼저 전파된 후 matmul와 같은 연산이 전파됩니다. Shardy에서는 사용자가 자체 op 우선순위를 설정할 수 있도록 허용하려고 합니다. Shardy에서는 언어를 사전에 알 수 없기 때문입니다. 따라서 Shardy에서 전파할 순서대로 작업 목록을 전달하도록 요청합니다.

아래 그림은 GSPMD에서 우선순위를 사용하여 올바른 순서로 작업을 전파하는 방법을 보여줍니다.

운영 우선순위가 중요한 이유에 관한 논의는 GSPMD 문서를 참고하세요.

방언에 관계없음

이전 인터페이스, 트레잇, 패스를 구현하는 한 Shardy는 언어에 맞게 작동할 수 있습니다. Shardy를 더 유연하고 방언에 구애받지 않도록 만들기 위해 노력하고 있으니 새로운 소식을 기대해 주세요.