<commit_before>local GenSGD,parent = torch.class('nn.GenSGDOptimization',
                                  'nn.BatchOptimization') 

-- this module parallelizes SGD in a particular way.  It sends out the
-- same batch to each of several workers, each with a different learning
-- rate.  The workers run and the parameters from the best worker and
-- it's learning rate are kept for the next batch.

function GenSGD:__init(...)
   parent.__init(self,...)
   xlua.unpack_class(self, {...},
                     'GenSGDOptimization', nil,
                     {arg='maxIterations', type='number', 
                      help='maximum nb of iterations per pass', default=1},
                     {arg='learningRate', type='number', 
                      help='learning rate (W = W - rate*dE/dW)', default=1e-2},
                     {arg='learningRateDecay', type='number', 
                      help='learning rate decay (lr_t = lr_0 / (1 + samplesSeen*lrDecay))', 
                      default=0},
                     {arg='weightDecay', type='number', 
                      help='amount of weight decay (W = W - decay*W)', default=0},
                     {arg='momentum', type='number', 
                      help='amount of momentum on weights (dE/W = dE/dW*(1-momentum) + prev(dE/dW)*momentum)', default=0},
                     {arg='sigma', type='number',
                      help='sigma of gaussian used to randomize learningRate',
                      default = 1e3}
                  )
   require 'lab' 
   if self.parallelize < 2 then
      xerror('GenSGD needs to work on several processors: set parallelize',
             'nn.GenSGDOptimization')
   end
   -- change the mapper to send the same batch to each worker
   self.copyBatch = true
   -- create default parameter set which will be randomized for each worker
   self.baseParameters = { momentum          = self.momentum, 
                           weightDecay       = self.weightDecay,
                           learningRate      = self.learningRate,
                           learningRateDecay = self.learningRateDecay,
                           sampleCounter     = self.sampleCounter
                        }
end

-- we are changing the way we map and reduce.  It would be nice to
-- change gradParametersPartial to ParametersPartial, as the logic is
-- different for this kind of parallelization.
function GenSGD:map_hook()
   local P = self.parallelize
   -- transmit new parameters to all workers
   self.children:join()
   self.children:send(self.parameters)
   print('randomizing for '..P..' lr: '..self.learningRate..' sigma: '..self.sigma)
   -- randomize learning rate (could randomize other bits)
   local n = torch.Tensor(P)

   n[1] = self.learningRate
   n[2] = self.learningRate * 10
   n[3] = self.learningRate / 10
   n[4] = self.learningRate / 100 
--  (lab.randn(P) * self.sigma):add(self.learningRate)
   self.baseParameters.sampleCounter = self.sampleCounter

   for t = 1,P do
      self.baseParameters.learningRate = n[t]
      print('lr: '..self.baseParameters.learningRate)
      --self.children[t]:join() 
      self.children[t]:send(self.baseParameters) 
   end
   -- then wait for all workers to return their Parameters + outputs
   -- should rename this to parametersParallel and optionsParallel
   gradParametersPartial = self.children:receive()
   outputsPartial = self.children:receive()
   -- force cleanup
   collectgarbage()
end

function GenSGD:reduce_hook()
   local P = self.parallelize
   local id = 0
   local mx = 1e9
   for t = 1,P do
      if outputsPartial[t].f_x < mx then
         id = t
         mx = outputsPartial[t].f_x
      end
   end
   if id == 0 then
      xerror('diverging','nn.GenSGDOptimization')
   else
      self.baseParameters = outputsPartial[id]
      self.output = self.baseParameters.f_x
      -- in this case we get the parameters back directly
      self.parameters:copy(gradParametersPartial[id])
      print('Winner: output = '..self.output..
            'learningRate = '..self.baseParameters['learningRate'])
      self.learningRate = self.baseParameters.learningRate
   end
end

function GenSGD:optimize()
   self.evaluate()
end

-- optimization (could do others in this mode)
GenSGD.optimizer = 
   function (module,params)
      -- apply momentum (store in the module)<commit_msg>log-normal dist for randome learning rates<commit_after>local GenSGD,parent = torch.class('nn.GenSGDOptimization',
                                  'nn.BatchOptimization') 

-- this module parallelizes SGD in a particular way.  It sends out the
-- same batch to each of several workers, each with a different learning
-- rate.  The workers run and the parameters from the best worker and
-- it's learning rate are kept for the next batch.

function GenSGD:__init(...)
   parent.__init(self,...)
   xlua.unpack_class(self, {...},
                     'GenSGDOptimization', nil,
                     {arg='maxIterations', type='number', 
                      help='maximum nb of iterations per pass', default=1},
                     {arg='learningRate', type='number', 
                      help='learning rate (W = W - rate*dE/dW)', default=1e-2},
                     {arg='learningRateDecay', type='number', 
                      help='learning rate decay (lr_t = lr_0 / (1 + samplesSeen*lrDecay))', 
                      default=0},
                     {arg='weightDecay', type='number', 
                      help='amount of weight decay (W = W - decay*W)', default=0},
                     {arg='momentum', type='number', 
                      help='amount of momentum on weights (dE/W = dE/dW*(1-momentum) + prev(dE/dW)*momentum)', default=0}
                  )
   require 'lab' 
   if self.parallelize < 2 then
      xerror('GenSGD needs to work on several processors: set parallelize',
             'nn.GenSGDOptimization')
   end
   -- change the mapper to send the same batch to each worker
   self.copyBatch = true
   -- create default parameter set which will be randomized for each worker
   self.baseParameters = { momentum          = self.momentum, 
                           weightDecay       = self.weightDecay,
                           learningRate      = self.learningRate,
                           learningRateDecay = self.learningRateDecay,
                           sampleCounter     = self.sampleCounter
                        }
end

-- we are changing the way we map and reduce.  It would be nice to
-- change gradParametersPartial to ParametersPartial, as the logic is
-- different for this kind of parallelization.
function GenSGD:map_hook()
   local P = self.parallelize
   -- transmit new parameters to all workers
   self.children:join()
   self.children:send(self.parameters)
   -- randomize learning rate (could randomize other bits).  Using a
   -- log normal around the base rate.
   local n = lab.randn(P):exp() * self.learningRate

   self.baseParameters.sampleCounter = self.sampleCounter

   for t = 1,P do
      self.baseParameters.learningRate = n[t]
      --self.children[t]:join() 
      self.children[t]:send(self.baseParameters) 
   end
   -- then wait for all workers to return their Parameters + outputs
   -- should rename this to parametersParallel and optionsParallel
   gradParametersPartial = self.children:receive()
   outputsPartial = self.children:receive()
   -- force cleanup
   collectgarbage()
end

function GenSGD:reduce_hook()
   local P = self.parallelize
   local id = 0
   local mx = 1e9
   for t = 1,P do
      if outputsPartial[t].f_x < mx then
         id = t
         mx = outputsPartial[t].f_x
      end
   end
   if id == 0 then
      xerror('diverging','nn.GenSGDOptimization')
   else
      self.baseParameters = outputsPartial[id]
      self.output = self.baseParameters.f_x
      -- in this case we get the parameters back directly
      self.parameters:copy(gradParametersPartial[id])
      -- keep this learning rate for the next batch
      self.learningRate = self.baseParameters.learningRate
   end
end

function GenSGD:optimize()
   self.evaluate()
end

-- optimization (could do others in this mode)
GenSGD.optimizer = 
   function (module,params)
      -- apply momentum (store in the module)
{'fraction_non_alphanumeric': 0.09869940440014587, 'fraction_numerical': 0.005469794578825817, 'mean_word_length': 2.1903838697169444, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4474018', 'n_tokens_mistral': 2319, 'n_tokens_neox': 2207, 'n_words': 683}