import numpy as np
import cupy
# cupy.cuda.set_allocator(cupy.cuda.MemoryPool(cupy.cuda.memory.malloc_managed).malloc)
import torch
import os
from torch.utils.dlpack import to_dlpack
from torch.utils.dlpack import from_dlpack
class PMemoryMM(cupy.cuda.memory.BaseMemory):
def __init__(self, size):
self.size = size
self.device_id = cupy.cuda.device.get_device_id()
self.ptr = 0
if size > 0:
self.ptr = cupy.cuda.runtime.hostAlloc(size, 0)
def __del__(self):
if self.ptr:
cupy.cuda.runtime.freeHost(self.ptr)
def my_pinned_allocatorMM(bsize):
return cupy.cuda.memory.MemoryPointer(PMemory(bsize),0)
class _CommonMM():
def _preInit(self):
fileNumber = 0
while os.path.isfile( self.diskname + str(fileNumber) + '.memmap.cpy.npy' ) == True:
fileNumber = fileNumber + 1
else:
self.fileName = self.diskname + str(fileNumber) + '.memmap'
numpyMemmap = np.memmap( self.fileName, dtype='float32', mode='w+', shape=(self.total_classes ,self.embed_dimension ))
np.save( self.fileName + '.cpy' , numpyMemmap, allow_pickle=True)
del numpyMemmap
os.remove(self.fileName)
def getCupyMM(self):
return self.CUPYmemmap
def saveCupy(self, saveFileName):
cupy.save( saveFileName, self.CUPYmemmap)
def getNumpyVersion(self):
return cupy.asnumpy(self.CUPYmemmap)
def _getReshapedRetrieval( self, retrievedPosIndexes , retrievedNegIndexes = None):
if not retrievedNegIndexes is None:
reshapedRetrieval = np.concatenate( [ retrievedPosIndexes.reshape(-1) , retrievedNegIndexes.reshape(-1) ] )
else:
reshapedRetrieval = retrievedPosIndexes.reshape(-1)
return reshapedRetrieval
class ModelFactoryMM(_CommonMM):
def __init__(self, model_variable, total_classes, embed_dimension, diskname = 'variable', datatype = 'float32', CPUPinn = False):
self.model_variable = model_variable
self.total_classes = total_classes
self.embed_dimension = embed_dimension
self.diskname = diskname
self.datatype = datatype
self.CPUPinn = CPUPinn
def zerosInit(self ):
#Initialize the memmap with just zeros
if self.CPUPinn == True:
cupy.cuda.set_allocator(my_pinned_allocator)
self._preInit()
self.CUPYmemmap = cupy.load( self.fileName+'.cpy.npy' , mmap_mode = 'r+' )
if self.CPUPinn == True:
cupy.cuda.set_allocator(None)
def uniformDistributionInit(self, low, high):
#Initialize the memmap with a uniform distribution
if self.CPUPinn == True:
cupy.cuda.set_allocator(my_pinned_allocator)
self._preInit()
self.CUPYmemmap = cupy.load( self.fileName+'.cpy.npy' , mmap_mode = 'r+' )
if self.total_classes > 100000:
for i in range( int( self.total_classes/100000) ):
j=i*100000
self.CUPYmemmap[j:j+100000] = cupy.random.uniform(low=low, high=high, size=(100000, self.embed_dimension), dtype=self.datatype)
for i in range( int( self.total_classes/100000)*100000, int( self.total_classes/10000) ):
j=i*10000
self.CUPYmemmap[j:j+10000] = cupy.random.uniform(low=low, high=high, size=(10000, self.embed_dimension), dtype=self.datatype)
for i in range( int( self.total_classes /10000)*10000 , self.total_classes ):
self.CUPYmemmap[i] = cupy.random.uniform(low=low, high=high, size=(self.embed_dimension), dtype=self.datatype)
elif self.total_classes > 10000:
for i in range( int(self.total_classes/10000) ):
j=i*10000
self.CUPYmemmap[j:j+10000] = cupy.random.uniform(low=low, high=high, size=(10000, self.embed_dimension), dtype=self.datatype)
for i in range( int( self.total_classes/10000)*10000 , self.total_classes ):
self.CUPYmemmap[i] = cupy.random.uniform(low=low, high=high, size=(self.embed_dimension), dtype=self.datatype)
else:
for i in range( self.total_classes ):
self.CUPYmemmap[i] = cupy.random.uniform(low=low, high=high, size=(self.embed_dimension), dtype=self.datatype)
if self.CPUPinn == True:
cupy.cuda.set_allocator(None)
def normalDistributionInit(self, mean, stdDev):
#Initialize the memmap with a normal distribution
if self.CPUPinn == True:
cupy.cuda.set_allocator(my_pinned_allocator)
self._preInit()
self.CUPYmemmap = cupy.load( self.fileName+'.cpy.npy' , mmap_mode = 'r+' )
if self.total_classes > 100000:
for i in range( int(self.total_classes/100000) ):
j=i*100000
self.CUPYmemmap[j:j+100000] = cupy.random.normal(loc=mean, scale=stdDev, size=(100000, self.embed_dimension), dtype=self.datatype )
for i in range( int(self.total_classes/100000)*100000, int(self.total_classes/10000) ):
j=i*10000
self.CUPYmemmap[j:j+10000] = cupy.random.normal(loc=mean, scale=stdDev, size=(10000, self.embed_dimension), dtype=self.datatype )
for i in range( int(self.total_classes/10000)*10000 , self.total_classes ):
self.CUPYmemmap[i] = cupy.random.normal(loc=mean, scale=stdDev, size=(self.embed_dimension), dtype=self.datatype )
elif self.total_classes > 10000:
for i in range( int(self.total_classes/10000) ):
j=i*10000
self.CUPYmemmap[j:j+10000] = cupy.random.normal(loc=mean, scale=stdDev, size=(10000, self.embed_dimension), dtype=self.datatype )
for i in range( int(self.total_classes/10000)*10000 , self.total_classes ):
self.CUPYmemmap[i] = cupy.random.normal(loc=mean, scale=stdDev, size=(self.embed_dimension), dtype=self.datatype )
else:
for i in range( self.total_classes ):
self.CUPYmemmap[i] = cupy.random.normal(loc=mean, scale=stdDev, size=(self.embed_dimension), dtype=self.datatype )
if self.CPUPinn == True:
cupy.cuda.set_allocator(None)
def variableTransformer(self, batchSize, posPerBatch, negPerBatch = None ):
if not negPerBatch == None:
return (np.arange( batchSize*posPerBatch ).reshape( batchSize , posPerBatch),
np.arange(start = batchSize*posPerBatch,
stop = batchSize*posPerBatch + batchSize*negPerBatch ).reshape(batchSize, negPerBatch) )
else:
return np.arange( batchSize*posPerBatch ).reshape( batchSize, posPerBatch )
def beforeForwardPass(self, retrievedPosIndexes , retrievedNegIndexes = None):
reshapedRetrieval = self._getReshapedRetrieval( retrievedPosIndexes, retrievedNegIndexes )
self.model_variable.weight.data = (
from_dlpack(self.CUPYmemmap[ reshapedRetrieval ].toDlpack() ) )
def afterOptimizerStep(self,retrievedPosIndexes , retrievedNegIndexes = None):
reshapedRetrieval = self._getReshapedRetrieval( retrievedPosIndexes, retrievedNegIndexes )
self.CUPYmemmap[ reshapedRetrieval ] = (
cupy.fromDlpack( to_dlpack( self.model_variable.weight.data ) ) )
class OptimizerFactoryMM(_CommonMM): #to do later, able to load matrixes to continue training
#take into account different size embedding matrices
def __init__(self, given_optimizer, total_classes, embed_dimension, model, variable_name, dtype='float32' , CPUPinn = False):
self.given_optimizer = given_optimizer
self.total_classes = total_classes
self.embed_dimension = embed_dimension
self.model = model
self.variable_name = variable_name
self.dtype = dtype
optimizer_index = None
self.CPUPinn = CPUPinn
#Some optiizers do not initialize its state until after first step
#So they need to initialized here
for group in given_optimizer.param_groups:
for p in group['params']:
state = given_optimizer.state[p]
# State initialization
if given_optimizer.__str__().split(' ', 1)[0] == 'SparseAdam':
# State initialization
if len(state) == 0:
state['step'] = 0
state['exp_avg'] = torch.zeros_like(p.data)
state['exp_avg_sq'] = torch.zeros_like(p.data)
self.optVarList = [ 'exp_avg', 'exp_avg_sq']
elif given_optimizer.__str__().split(' ', 1)[0] == 'Adagrad':
self.optVarList = [ 'sum' ]
elif given_optimizer.__str__().split(' ', 1)[0] == 'Adadelta':
if len(state) == 0:
state['step'] = 0
state['square_avg'] = torch.zeros_like(p.data)
state['acc_delta'] = torch.zeros_like(p.data)
self.optVarList = [ 'square_avg', 'acc_delta']
elif given_optimizer.__str__().split(' ', 1)[0] == 'Adamax':
if len(state) == 0:
state['step'] = 0
state['exp_avg'] = torch.zeros_like(p.data)
state['exp_inf'] = torch.zeros_like(p.data)
self.optVarList = [ 'exp_avg', 'exp_inf']
elif given_optimizer.__str__().split(' ', 1)[0] == 'RMSprop':
if len(state) == 0:
state['step'] = 0
state['square_avg'] = torch.zeros_like(p.data)
if group['momentum'] > 0:
state['momentum_buffer'] = torch.zeros_like(p.data)
if group['centered']:
state['grad_avg'] = torch.zeros_like(p.data)
self.optVarList = [ 'square_avg']
if group['momentum'] > 0:
self.optVarList.append( 'momentum_buffer' )
if group['centered']:
self.optVarList.append( 'grad_avg' )
elif given_optimizer.__str__().split(' ', 1)[0] == 'Rprop':
if p.grad is None:
print('Error, gradients are empty')
print('For Rprop, need to first run at least 1 training step that has gradients')
return
if len(state) == 0:
state['step'] = 0
state['prev'] = torch.zeros_like(p.data)
#For now, do now know how to Not initialize this due to len(state)==0 in optimizer
state['step_size'] = grad.new().resize_as_(grad).fill_(group['lr'])
self.optVarList = [ 'prev']
elif given_optimizer.__str__().split(' ', 1)[0] == 'ASGD':
if len(state) == 0:
state['step'] = 0
state['eta'] = group['lr']
state['mu'] = 1
state['ax'] = torch.zeros_like(p.data)
self.optVarList = [ 'ax']
elif given_optimizer.__str__().split(' ', 1)[0] == 'AdamW':
amsgrad = group['amsgrad']
if len(state) == 0:
state['step'] = 0
state['exp_avg'] = torch.zeros_like(p.data)
state['exp_avg_sq'] = torch.zeros_like(p.data)
if amsgrad:
state['max_exp_avg_sq'] = torch.zeros_like(p.data)
self.optVarList = [ 'exp_avg', 'exp_avg_sq']
if amsgrad:
self.optVarList.append('max_exp_avg_sq')
elif given_optimizer.__str__().split(' ', 1)[0] == 'Adam':
amsgrad = group['amsgrad']
if len(state) == 0:
state['step'] = 0
state['exp_avg'] = torch.zeros_like(p.data)
state['exp_avg_sq'] = torch.zeros_like(p.data)
if amsgrad:
state['max_exp_avg_sq'] = torch.zeros_like(p.data)
self.optVarList = [ 'exp_avg', 'exp_avg_sq']
if amsgrad:
self.optVarList.append( 'max_exp_avg_sq' )
else:
print('This optimizer is not currently supported. Please choose a different optimizer')
return
#Figure out which index for given variable
for i, item in enumerate( self.model.named_parameters() ):
if item[0][:-7] == self.variable_name:
optimizer_index = i
self.diskname = item[0][:-7] + given_optimizer.__str__().split(' ', 1)[0]
if optimizer_index == None:
print( 'Error: No variable with that name is in Model. Please initialize again with correct name' )
return
optimizerKeyList = list(self.given_optimizer.state_dict()['state'].keys())
self.optimizerKey = optimizerKeyList[ optimizer_index ]
def _preInit(self):
for optVar in self.optVarList:
fileNumber = 0
while os.path.isfile( self.diskname + str(fileNumber) + '.memmap.cpy.npy' ) == True:
fileNumber = fileNumber + 1
else:
self.fileName = self.diskname + str(fileNumber) + '.memmap'
numpyMemmap = np.memmap( self.fileName+optVar, dtype='float32', mode='w+', shape=(self.total_classes ,self.embed_dimension ))
np.save( self.fileName + optVar + '.cpy' , numpyMemmap, allow_pickle=True)
del numpyMemmap
os.remove(self.fileName+optVar)
def optInit(self):
if self.CPUPinn == True:
cupy.cuda.set_allocator(my_pinned_allocator)
self._preInit()
self.CUPYmemmap = []
for optVar in self.optVarList:
self.CUPYmemmap.append( cupy.load( self.fileName+optVar+'.cpy.npy' , mmap_mode = 'r+' ) )
if self.CPUPinn == True:
cupy.cuda.set_allocator(None)
def beforeForwardPass(self, retrievedPosIndexes , retrievedNegIndexes = None):
reshapedRetrieval = self._getReshapedRetrieval( retrievedPosIndexes, retrievedNegIndexes )
for idx, optVar in enumerate(self.optVarList):
self.given_optimizer.state_dict()['state'][ self.optimizerKey ][optVar] = (
from_dlpack( self.CUPYmemmap[idx][ reshapedRetrieval ].toDlpack() ) )
def afterOptimizerStep(self, retrievedPosIndexes , retrievedNegIndexes = None):
reshapedRetrieval = self._getReshapedRetrieval( retrievedPosIndexes, retrievedNegIndexes )
for idx, optVar in enumerate(self.optVarList):
self.CUPYmemmap[idx][ reshapedRetrieval ] = (
cupy.fromDlpack( to_dlpack( self.given_optimizer.state_dict()['state'][ self.optimizerKey ][optVar] ) ) )
class COMMM(_CommonMM):
def __init__(self, total_classes, diskname = 'COM', datatype = 'uint32', CPUPinn = False ):
self.total_classes = total_classes
self.datatype = datatype
self.diskname = diskname
self.CPUPinn = CPUPinn
def _preInit(self): #Can't depend on inherited since the shape is different
fileNumber = 0
while os.path.isfile(diskself.disknamename+str(fileNumber) + 'memmap' ) == false:
fileNumber = fileNumber + 1
else:
self.fileName = self.diskname+str(fileNumber)
numpyMemmap = np.memmap(self.fileName, dtype=datatype, mode='w+', shape=(total_classes , total_classes ))
np.save( self.fileName+'.cpy' , numpyMemmap, allow_pickle=True)
del numpyMemmap
os.remove(self.fileName)
def comInit(self, CPUPinn=False):
if self.CPUPinn == True:
cupy.cuda.set_allocator(my_pinned_allocator)
self._preInit()
self.CUPYmemmap = cupy.load( fileName+'.cpy.npy' , mmap_mode = 'r+' )
if self.CPUPinn == True:
cupy.cuda.set_allocator(None)
class DataGadgetMM(_CommonMM):
def __init__(self, fileName, CPUPinn=False):
self.Numpyfilename = Numpyfilename
self.CPUPinn = CPUPinn
def gadgetInit(self):
if self.CPUPinn == True:
cupy.cuda.set_allocator(my_pinned_allocator)
self.CUPYmemmap = cupy.load( self.fileName , mmap_mode = 'r+' )
if self.CPUPinn == True:
cupy.cuda.set_allocator(None)
def getData(self, retrievedPosIndexes , retrievedNegIndexes = None):
return from_dlpack( self.CUPYmemmap[ reshapedRetrieval ].toDlpack() )
