From b49d6b69936c5898007dda9fcf95fdfb2aa8ce4d Mon Sep 17 00:00:00 2001 From: zhoushenglong Date: Fri, 29 Nov 2024 08:33:44 +0000 Subject: [PATCH 1/2] add env to support different mm layout on maca. --- lmdeploy/pytorch/backends/dlinfer/linear.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/lmdeploy/pytorch/backends/dlinfer/linear.py b/lmdeploy/pytorch/backends/dlinfer/linear.py index 567a01ddd..f2450a974 100644 --- a/lmdeploy/pytorch/backends/dlinfer/linear.py +++ b/lmdeploy/pytorch/backends/dlinfer/linear.py @@ -1,4 +1,5 @@ # Copyright (c) OpenMMLab. All rights reserved. +import os from typing import Optional import torch @@ -11,6 +12,14 @@ class DlinferLinearImpl(LinearImpl): """Dlinfer linear implementation api.""" + def update_weights(self, + weight: torch.Tensor, + bias: Optional[torch.Tensor] = None): + """update weights.""" + if os.getenv('TORCH_MACA_NN_LAYOUT', 'False').lower() == 'true': + weight = weight.data.t().contiguous() + return weight, bias + def forward(self, x, weight: torch.Tensor, From c8362423feb1f6c423a1c401ce07b0c414bc0334 Mon Sep 17 00:00:00 2001 From: zhoushenglong Date: Fri, 29 Nov 2024 10:35:37 +0000 Subject: [PATCH 2/2] rename env variable. --- lmdeploy/pytorch/backends/dlinfer/linear.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lmdeploy/pytorch/backends/dlinfer/linear.py b/lmdeploy/pytorch/backends/dlinfer/linear.py index f2450a974..5edfa7728 100644 --- a/lmdeploy/pytorch/backends/dlinfer/linear.py +++ b/lmdeploy/pytorch/backends/dlinfer/linear.py @@ -16,7 +16,7 @@ def update_weights(self, weight: torch.Tensor, bias: Optional[torch.Tensor] = None): """update weights.""" - if os.getenv('TORCH_MACA_NN_LAYOUT', 'False').lower() == 'true': + if os.getenv('DLINER_LINEAR_USE_NN_LAYOUT', '0') == '1': weight = weight.data.t().contiguous() return weight, bias