# Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from functools import reduce from paddle.framework import core from paddle.framework.io import Variable dtype_to_size = { core.VarDesc.VarType.FP16: 2, core.VarDesc.VarType.FP32: 4, core.VarDesc.VarType.FP64: 8, core.VarDesc.VarType.INT16: 2, core.VarDesc.VarType.INT32: 4, core.VarDesc.VarType.INT64: 8, core.VarDesc.VarType.BOOL: 1, core.VarDesc.VarType.UINT8: 1, } class VarBlock: def __init__(self, varname, offset, size): self.varname = varname # NOTE: real offset is offset * size self.offset = offset self.size = size def __str__(self): return f"{self.varname}:{int(self.offset)}:{int(self.size)}" def create_var_struct(var): if var.type == core.VarDesc.VarType.SELECTED_ROWS: lod_level = None elif var.type == core.VarDesc.VarType.DENSE_TENSOR: lod_level = var.lod_level else: raise ValueError("can only support SELECTED_ROWS/DENSE_TENSOR now") return VarStruct( var.name, var.shape, var.dtype, var.type, lod_level, var.persistable ) class VarStruct: """ record part properties of a Variable in python. """ def __init__(self, name, shape, dtype, type, lod_level, persistable): self.name = name self.shape = shape self.dtype = dtype self.type = type self.lod_level = lod_level self.persistable = persistable self.m_size = 1 self.m_size = reduce(lambda x, y: x * y, shape, 1) self.m_size *= dtype_to_size[dtype] def __str__(self): return f"N: {self.name}, S: {self.shape}, D: {self.dtype}, T: {self.type}, LL: {self.lod_level}, P: {self.persistable}, M: {self.m_size}" class VarDistributed: """ a class to record the var distributed on parameter servers. the class will record the relationship between origin var and slice var. the slice var's properties, such as type/shape/offset/endpoint. """ def __init__( self, origin_var, slice_var, is_slice=None, block_id=None, offset=None, vtype=None, endpoint=None, ): """ Args: origin_var(Variable|VarStruct): origin var properties slice_var(Variable|VarStruct): slice var properties is_slice(bool|None): slice or not, slice_var=True/False and its block size > 8192 are the judgement standard. block_id(int|None): the number about the slice var. offset(int|None): if the slice var is sliced, offset is the numel before the var. vtype(str|None): a tag, such as Optimizer/Param/RemotePrefetch. endpoint(str|None): which parameter the slice var on, such as "127.0.0.1:1001" """ if isinstance(origin_var, Variable): self.origin = create_var_struct(origin_var) else: self.origin = origin_var if isinstance(slice_var, Variable): self.slice = create_var_struct(slice_var) else: self.slice = slice_var if self.equal(self.origin, self.slice): self.is_slice = False self.block_id = 0 self.offset = 0 else: self.is_slice = True self.block_id = 0 self.offset = 0 if is_slice is not None: self.is_slice = is_slice if block_id is not None: self.block_id = block_id if offset is not None: self.offset = offset self.vtype = vtype self.endpoint = endpoint @staticmethod def equal(var1, var2): """ the two var is equal or not. Returns: bool: equal will return True else False """ assert isinstance(var1, VarStruct) and isinstance(var2, VarStruct) return ( var1.name == var2.name and var1.type == var2.type and var1.shape == var2.shape and var1.dtype == var2.dtype and var1.lod_level == var2.lod_level and var1.persistable == var2.persistable ) def __str__(self): origin_var_str = f"{self.origin.name} : base.{self.origin.type}.shape{self.origin.shape}.astype({self.origin.dtype})" slice_var_str = ( f"{self.slice.name} : base.{self.slice.type}.shape{self.slice.shape}.astype({self.slice.dtype})" f".slice({self.is_slice}).block({self.block_id}).offset({self.offset})" ) return f"var owned: {self.vtype}, origin var: ( {origin_var_str} ), slice var: ( {slice_var_str} ), endpoint: {self.endpoint} " class VarsDistributed: """ a gather about VarDistributed with many methods to find distributed vars. through the class, we can get overview about the distributed parameters on parameter servers. this class may centralized and convenient for developer to manage and get variable's distribute. other module can also use this to find variables such io.py. """ def __init__(self): self.distributed_vars = [] def add_distributed_var( self, origin_var, slice_var, is_slice=None, block_id=None, offset=None, vtype=None, endpoint=None, ): """ add distributed var in this. Args: origin_var(Variable|VarStruct): origin var properties slice_var(Variable|VarStruct): slice var properties is_slice(bool|None): slice or not, slice_var=True/False and its block size > 8192 are the judgement standard. block_id(int|None): the number about the slice var. offset(int|None): if the slice var is sliced, offset is the numel before the var. vtype(str|None): a tag, such as Optimizer/Param/RemotePrefetch. endpoint(str|None): which parameter the slice var on, such as "127.0.0.1:1001" Returns: None """ self.distributed_vars.append( VarDistributed( origin_var, slice_var, is_slice, block_id, offset, vtype, endpoint, ) )