cleanups

Author: ngc92

src/training/gradients.cpp

@@ -5,7 +5,6 @@
 #include "gradients.h"
 
 #include "model.h"
-#include "models/llama_weights.h"
 #include "utilities/allocator.h"
 #include "utilities/comm.h"
 #include "utilities/lazy_allocator.h"
@@ -32,13 +31,13 @@ UnshardedGradientManager::UnshardedGradientManager(const TransformerConfig& cfg,
         mBlockGradients[i] = model.create_block_container(cfg, cfg.DType, cfg.DType);
         alloc_lazy.allocate(mBlockGradients[i]);
         alloc_lazy.commit(*alloc, EAllocationType::ON_DEVICE, "block_grad");
-        mBlockShards[i] = shard_view(GenericTensorContainer(mBlockGradients[i]), rank, world);
+        mBlockShards[i] = shard_view(mBlockGradients[i], rank, world);
     }
 
     mNonBlockGradients = model.create_non_block_container(cfg, cfg.DType, cfg.DType);
     alloc_lazy.allocate(mNonBlockGradients);
     alloc_lazy.commit(*alloc, EAllocationType::ON_DEVICE, "nonblock_grad");
-    mNonBlockShards = shard_view(GenericTensorContainer(mNonBlockGradients), rank, world);
+    mNonBlockShards = shard_view(mNonBlockGradients, rank, world);
 
     mGradEvent = create_named_event("grad_event");
 }
@@ -97,7 +96,7 @@ ShardedBlocksGradientManager::ShardedBlocksGradientManager(const TransformerConf
     mFullNonBlock = model.create_non_block_container(cfg, cfg.DType, cfg.DType);
     alloc_lazy.allocate(mFullNonBlock);
     alloc_lazy.commit(*alloc, EAllocationType::ON_DEVICE, "nonblock_grad");
-    mNonBlockShards = shard_view(GenericTensorContainer(mFullNonBlock), rank, world);
+    mNonBlockShards = shard_view(mFullNonBlock, rank, world);
 
     mGradBuffers[0] = model.create_block_container(cfg, cfg.DType, cfg.DType);
     mGradBuffers[1] = model.create_block_container(cfg, cfg.DType, cfg.DType);

src/utilities/tensor.cpp

@@ -160,17 +160,17 @@ const Tensor& GenericTensorContainer::get_tensor(std::size_t idx) const {
 
 GenericTensorContainer shard_empty_container(GenericTensorContainer&& c, int world) {
     // can't use visit here, because we explicitly want to iterate over empty tensors
-    for (int i = 0; i < c.num_tensors(); ++i) {
+    for (std::size_t i = 0; i < c.num_tensors(); ++i) {
         auto& t = c.get_tensor(i);
-        if (!t.empty()) { throw std::logic_error("shard_container called with non-empty tensor"); }
+        if (!t.empty()) { throw std::logic_error("shard_empty_container called with non-empty tensor"); }
         t.Sizes[0] = div_exact(t.Sizes[0], static_cast<long>(world));
     }
-    return c;
+    return std::move(c);
 }
 
 GenericTensorContainer shard_view(const GenericTensorContainer& c, int rank, int world) {
     std::vector<Tensor> shards(c.num_tensors());
-    for (int i = 0; i < c.num_tensors(); ++i) {
+    for (std::size_t i = 0; i < c.num_tensors(); ++i) {
         shards.at(i) = static_cast<Tensor>(shard_view(c.get_tensor(i), rank, world));
     }
     return GenericTensorContainer{shards};

src/utilities/tensor_container.h

@@ -57,11 +57,11 @@ class GenericTensorContainer final : public SimpleTensorContainer {
     std::vector<Tensor> mTensors;
 };
 
-//! shards an _empty_ container, i.e., a container in which all data pointers
+//! Shards an _empty_ container, i.e., a container in which all data pointers
 //! are `nullptr`, but sizes have been set up.
 GenericTensorContainer shard_empty_container(GenericTensorContainer&& c, int world);
 
-//! shard a non-empty tensor container. The returned container's tensors are _views_ into
+//! Shards a non-empty tensor container. The returned container's tensors are _views_ into
 //! the original container's tensors.
 GenericTensorContainer shard_view(const GenericTensorContainer& c, int rank, int world);