diff --git a/CMakeLists.txt b/CMakeLists.txt
index bac1314..a52fc3f 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -151,6 +151,7 @@ add_library(llmq-common SHARED
         src/training/logging.cpp
         src/training/checkpoint.cpp
         src/training/model.cpp
+        src/training/gradients.cpp
         src/training/transformer_config.cpp
 
         src/models/llama_run_state.cpp
diff --git a/src/binding/py_train.cpp b/src/binding/py_train.cpp
index 215031c..5f875d3 100644
--- a/src/binding/py_train.cpp
+++ b/src/binding/py_train.cpp
@@ -276,18 +276,21 @@ std::vector<std::pair<std::string, long>> MultiGPUPyTrainer::get_stack_info(int
 }
 
 std::vector<std::pair<std::string, Tensor>> MultiGPUPyTrainer::get_gradients(int gpu_id) {
+    using namespace LLamaWeightID;
+
     std::vector<std::pair<std::string, Tensor>> result;
+    // TODO make this work with generalized gradients
     run_work([&result](sThreadContext& ctx) {
         const auto& config = ctx.Model->config();
         auto& grads = ctx.Model->grads();
         CUDA_CHECK(cudaDeviceSynchronize());
-        result.emplace_back("model.embed_tokens.weight", grads.get_embeddings_shard(nullptr));
+        result.emplace_back("model.embed_tokens.weight", grads.get_non_block_shard(LLamaWeightID::EMBEDDING, nullptr));
         if (!config.TiedWordEmbeddings) {
-            result.emplace_back("lm_head.weight", grads.get_lmhead_shard(nullptr));
+            result.emplace_back("lm_head.weight", grads.get_non_block_shard(LM_HEAD, nullptr));
         }
-        result.emplace_back("model.norm.weight", grads.get_lnf_w_shard(nullptr));
+        result.emplace_back("model.norm.weight", grads.get_non_block_shard(LNF_W, nullptr));
         for (int l = 0; l < config.NumLayers; l++) {
-            using namespace LLamaWeightID;
+
             std::string prefix = "model.layers." + std::to_string(l);
             auto& block = grads.get_block_shard(l, nullptr);
             result.emplace_back(prefix + ".self_attn.qkv.weight", block.get_tensor(QKV_W));
diff --git a/src/models/llama_gradients.cpp b/src/models/llama_gradients.cpp
index 3957a86..d73f362 100644
--- a/src/models/llama_gradients.cpp
+++ b/src/models/llama_gradients.cpp
@@ -8,339 +8,49 @@
 #include "llama_model.h"
 #include "utilities/comm.h"
 
-LLamaGradsManager::LLamaGradsManager(std::uint64_t seed, int step) : mRng(seed), mStepCounter(step) {
-
-}
-
-void LLamaGradsManager::scatter_reduce(Tensor& tensor, cudaStream_t stream, cudaEvent_t signal, NCCLCommunicator& comm) {
-    comm.begin_transaction(stream);
-    comm.schedule_reduce_scatter(tensor);
-    comm.execute_transaction(signal);
-}
-
-void LLamaGradsManager::scatter_reduce(int layer_idx, SimpleTensorContainer& block, cudaStream_t stream, cudaEvent_t signal, NCCLCommunicator& comm) {
-    comm.begin_transaction(stream);
-    visit([&](Tensor& t){ comm.schedule_reduce_scatter(t); }, block);
-    comm.execute_transaction(signal);
-}
-
-void LLamaGradsManager::start_micro_step(cudaStream_t stream, int micro_step, int total_steps) {
-    mIsFirstMicroStep = micro_step == 0;
-    mIsLastMicroStep = micro_step == total_steps - 1;
-    if (micro_step == 0) {
-        ++mStepCounter;
-        on_first_micro_step(stream);
-    }
-}
-
-class LLamaGradientsUnsharded : public LLamaGradsManager {
-public:
-    LLamaGradientsUnsharded(std::uint64_t seed, int step, const TransformerConfig& config, int rank, int world, const std::shared_ptr<TensorAllocator>& alloc);
-
-    void on_first_micro_step(cudaStream_t stream) override;
-    void end_micro_step(cudaStream_t stream, NCCLCommunicator& comm) override;
-
-    Tensor& get_embeddings_full(cudaStream_t stream, NCCLCommunicator& comm, bool& accumulate) override;
-    Tensor& get_lmhead_full(cudaStream_t stream, NCCLCommunicator& comm, bool& accumulate) override;
-    Tensor& get_lnf_w_full(cudaStream_t stream, NCCLCommunicator& comm, bool& accumulate) override;
-    sLLamaBlockWeights<Tensor>& get_block_full(int layer_idx, cudaStream_t stream, NCCLCommunicator& comm, bool& accumulate) override;
-
-    TensorShard& get_embeddings_shard(cudaStream_t stream) override;
-    TensorShard& get_lmhead_shard(cudaStream_t stream) override;
-    TensorShard& get_lnf_w_shard(cudaStream_t stream) override;
-    sLLamaBlockWeights<TensorShard>& get_block_shard(int layer_idx, cudaStream_t stream) override;
-
-    void notify_embeddings(cudaStream_t stream, NCCLCommunicator& comm) override;
-    void notify_lmhead(cudaStream_t stream, NCCLCommunicator& comm) override;
-    void notify_lnf_w(cudaStream_t stream, NCCLCommunicator& comm) override;
-
-    void notify_block(int layer_idx, cudaStream_t stream, NCCLCommunicator& comm) override;
-private:
-    sLLamaWeightsSet<Tensor> mFullGradient;
-    sLLamaWeightsSet<TensorShard> mShardView;
-    cudaEvent_t mGradEvent;
-};
-
-LLamaGradientsUnsharded::LLamaGradientsUnsharded(std::uint64_t seed, int step, const TransformerConfig& config, int rank, int world, const std::shared_ptr<TensorAllocator>& alloc) :
-    LLamaGradsManager(seed, step)
-{
-    mFullGradient = allocate_full_weights(config, EAllocationType::ON_DEVICE, *alloc);
-    mShardView.NonBlocks = shard_non_block(mFullGradient.NonBlocks, rank, world);
-    mShardView.Blocks.reserve(config.NumLayers);
-    for(int i = 0; i < config.NumLayers; ++i) {
-        mShardView.Blocks.push_back(shard_block(mFullGradient.Blocks[i], rank, world));
-    }
-    mGradEvent = create_named_event("grad_event");
-}
-
 void LLamaGradientsUnsharded::on_first_micro_step(cudaStream_t stream) {
-    fill_zero(mFullGradient.NonBlocks.LNF_w, stream);
-    if(mFullGradient.NonBlocks.LMHead.Data != mFullGradient.NonBlocks.Embeddings.Data) {
-        fill_zero(mFullGradient.NonBlocks.Embeddings, stream);
-        fill_zero(mFullGradient.NonBlocks.LMHead, stream);
+    using namespace LLamaWeightID;
+    fill_zero(mNonBlockGradients.get_tensor(LNF_W), stream);
+    if(mNonBlockGradients.get_tensor(LM_HEAD).Data != nullptr) {
+        fill_zero(mNonBlockGradients.get_tensor(EMBEDDING), stream);
     } else {
         // embedding backward comes after LMHead backward; and LMHead backward *sets* the gradient
         // on the first backward call, so no need to zero anything.
     }
-    for(auto& layer: mFullGradient.Blocks) {
-        fill_zero(layer.LN1_w, stream);
-        fill_zero(layer.LN2_w, stream);
-        fill_zero(layer.Attn_QKV_b, stream);
+    for(auto& layer: mBlockGradients) {
+        fill_zero(layer.get_tensor(LN1_W), stream);
+        fill_zero(layer.get_tensor(LN2_W), stream);
+        fill_zero(layer.get_tensor(QKV_B), stream);
         // no need to zero out the matrix weights, we'll just overwrite them on the first
         // grad accumulation step
     }
 }
 
-void LLamaGradientsUnsharded::end_micro_step(cudaStream_t stream, NCCLCommunicator& comm) {
-    if (mIsLastMicroStep) {
-        CUDA_CHECK(cudaStreamWaitEvent(stream, mGradEvent, 0));
-    }
-}
-
-
-Tensor& LLamaGradientsUnsharded::get_embeddings_full(cudaStream_t stream, NCCLCommunicator& comm, bool& accumulate) {
-    accumulate = !mIsFirstMicroStep;
-    return mFullGradient.NonBlocks.Embeddings;
-}
-Tensor& LLamaGradientsUnsharded::get_lmhead_full(cudaStream_t stream, NCCLCommunicator& comm, bool& accumulate) {
-    accumulate = !mIsFirstMicroStep;
-    return mFullGradient.NonBlocks.LMHead;
-}
-Tensor& LLamaGradientsUnsharded::get_lnf_w_full(cudaStream_t stream, NCCLCommunicator& comm, bool& accumulate) {
-    accumulate = !mIsFirstMicroStep;
-    return mFullGradient.NonBlocks.LNF_w;
-}
-sLLamaBlockWeights<Tensor>& LLamaGradientsUnsharded::get_block_full(int layer_idx, cudaStream_t stream, NCCLCommunicator& comm, bool& accumulate) {
-    accumulate = !mIsFirstMicroStep;
-    return mFullGradient.Blocks.at(layer_idx);
-}
-
-TensorShard& LLamaGradientsUnsharded::get_embeddings_shard(cudaStream_t stream) {
-    return mShardView.NonBlocks.Embeddings;
-}
-TensorShard& LLamaGradientsUnsharded::get_lmhead_shard(cudaStream_t stream) {
-    return mShardView.NonBlocks.LMHead;
-}
-TensorShard& LLamaGradientsUnsharded::get_lnf_w_shard(cudaStream_t stream) {
-    return mShardView.NonBlocks.LNF_w;
-}
-sLLamaBlockWeights<TensorShard>& LLamaGradientsUnsharded::get_block_shard(int layer_idx, cudaStream_t stream) {
-    return mShardView.Blocks.at(layer_idx);
-}
-
-void LLamaGradientsUnsharded::notify_embeddings(cudaStream_t stream, NCCLCommunicator& comm) {
-    if(!mIsLastMicroStep) return;
-    if (comm.world_size() != 1) {
-        NvtxRange r{"notify_embeddings"};
-        scatter_reduce(mFullGradient.NonBlocks.Embeddings, stream, mGradEvent, comm);
-    }
-}
-
-void LLamaGradientsUnsharded::notify_lmhead(cudaStream_t stream, NCCLCommunicator& comm) {
-    if(!mIsLastMicroStep) return;
-    if(mFullGradient.NonBlocks.LMHead.Data == mFullGradient.NonBlocks.Embeddings.Data) return;    // sync lmhead with embeddings
-    if (comm.world_size() != 1) {
-        NvtxRange r{"notify_lmhead"};
-        scatter_reduce(mFullGradient.NonBlocks.LMHead, stream, mGradEvent, comm);
-    }
-}
-
-void LLamaGradientsUnsharded::notify_lnf_w(cudaStream_t stream, NCCLCommunicator& comm) {
-    if(!mIsLastMicroStep) return;
-    if (comm.world_size() != 1) {
-        scatter_reduce(mFullGradient.NonBlocks.LNF_w, stream, mGradEvent, comm);
-    }
-}
-
-void LLamaGradientsUnsharded::notify_block(int layer_idx, cudaStream_t stream, NCCLCommunicator& comm) {
-    if(!mIsLastMicroStep) return;
-    auto& dw = mFullGradient.Blocks[layer_idx];
-    if (comm.world_size() != 1) {
-        scatter_reduce(layer_idx, dw, stream, mGradEvent, comm);
-    }
-}
-
 // ---------------------------------------------------------------------------------------------------------------------
 
 // shard the transformer blocks, but not the embeddings and lmhead.
-class LLamaGradientsBlockShardedBase : public LLamaGradsManager {
-public:
-    LLamaGradientsBlockShardedBase(std::uint64_t seed, int step, const TransformerConfig& config, const LLamaOptions& options, int rank, int world, const std::shared_ptr<TensorAllocator>& alloc);
-
-    void on_first_micro_step(cudaStream_t stream) override;
-    void end_micro_step(cudaStream_t stream, NCCLCommunicator& comm) override;
-
-    Tensor& get_embeddings_full(cudaStream_t stream, NCCLCommunicator& comm, bool& accumulate) override;
-    Tensor& get_lmhead_full(cudaStream_t stream, NCCLCommunicator& comm, bool& accumulate) override;
-    Tensor& get_lnf_w_full(cudaStream_t stream, NCCLCommunicator& comm, bool& accumulate) override;
-    sLLamaBlockWeights<Tensor>& get_block_full(int layer_idx, cudaStream_t stream, NCCLCommunicator& comm, bool& accumulate) override;
-
-    void notify_embeddings(cudaStream_t stream, NCCLCommunicator& comm) override;
-    void notify_lmhead(cudaStream_t stream, NCCLCommunicator& comm) override;
-    void notify_lnf_w(cudaStream_t stream, NCCLCommunicator& comm) override;
-    void notify_block(int layer_idx, cudaStream_t stream, NCCLCommunicator& comm) override;
-
-    TensorShard& get_embeddings_shard(cudaStream_t stream) override;
-    TensorShard& get_lmhead_shard(cudaStream_t stream) override;
-    TensorShard& get_lnf_w_shard(cudaStream_t stream) override;
-    sLLamaBlockWeights<TensorShard>& get_block_shard(int layer_idx, cudaStream_t stream) override;
-private:
-    virtual void sr_accumulate_layer(int layer_idx,
-                                     SimpleTensorContainer& dw,
-                                     SimpleTensorContainer& sw,
-                                     cudaStream_t stream,
-                                     NCCLCommunicator& comm) = 0;
-
-    sLLamaNonBlockWeights<Tensor> mFullNonBlock;
-    sLLamaNonBlockWeights<TensorShard> mNonBlockShards;
-
-    std::array<sLLamaBlockWeights<Tensor>, 2> mGradBuffers;
-    struct sBlockState {
-        cudaEvent_t Event;
-        int LayerIdx = -1;
-        bool NeedsAccumulation = false;
-    };
-    std::array<sBlockState, 2> mGradStates;
-    std::vector<sLLamaBlockWeights<TensorShard>> mGradShards;
-    cudaEvent_t mNonBlockEvent;
-};
-
-LLamaGradientsBlockShardedBase::LLamaGradientsBlockShardedBase(std::uint64_t seed, int step, const TransformerConfig& config, const LLamaOptions& options, int rank, int world, const std::shared_ptr<TensorAllocator>& alloc):
-    LLamaGradsManager(seed, step),
-    mFullNonBlock( allocate_non_block_full(config, config.DType, EAllocationType::ON_DEVICE, *alloc))
-{
-    mGradBuffers[0] = allocate_block_full(config, config.DType, config.DType, EAllocationType::ON_DEVICE, *alloc);
-    mGradBuffers[1] = allocate_block_full(config, config.DType, config.DType, EAllocationType::ON_DEVICE, *alloc);
-    mGradStates[0].Event = create_named_event("grad_event_0");
-    mGradStates[1].Event = create_named_event("grad_event_1");
-    mNonBlockEvent = create_named_event("grad_nonblock_event");
-    mGradShards.reserve(config.NumLayers);
-    for(int i = 0; i < config.NumLayers; ++i) {
-        EAllocationType kind = options.OffloadGrads ? EAllocationType::PINNED : EAllocationType::ON_DEVICE;
-        mGradShards.push_back(allocate_block_shard(config, config.DType, config.DType, kind, rank, world, *alloc));
-    }
-
-    mNonBlockShards = shard_non_block(mFullNonBlock, rank, world);
-}
 
 void LLamaGradientsBlockShardedBase::on_first_micro_step(cudaStream_t stream) {
-    // if we have untied embeddings, we need to zero them out
-    if(mFullNonBlock.Embeddings.Data != mFullNonBlock.LMHead.Data) {
-        fill_zero(mFullNonBlock.Embeddings, stream);
-    }
-    fill_zero(mFullNonBlock.LNF_w, stream);
-}
-
-void LLamaGradientsBlockShardedBase::end_micro_step(cudaStream_t stream, NCCLCommunicator& comm) {
-    for (int i = 0; i < 2; ++i) {
-        auto& state = mGradStates[i];
-        int layer_idx = state.LayerIdx;
-        if (state.NeedsAccumulation) {
-            // we need to wait for the previous accumulation to finish
-            CUDA_CHECK(cudaStreamWaitEvent(stream, state.Event, 0));
-            auto& dw = mGradBuffers.at(layer_idx % 2);
-            auto& sw = mGradShards.at(layer_idx);
-            sr_accumulate_layer(layer_idx, dw, sw, stream, comm);
-            state.NeedsAccumulation = false;
-        }
-    }
-    if (mIsLastMicroStep)
-        CUDA_CHECK(cudaStreamWaitEvent(stream, mNonBlockEvent, 0));
-}
-
-Tensor& LLamaGradientsBlockShardedBase::get_embeddings_full(cudaStream_t stream, NCCLCommunicator& comm, bool& accumulate) {
-    accumulate = !mIsFirstMicroStep;
-    return mFullNonBlock.Embeddings;
-}
-
-Tensor& LLamaGradientsBlockShardedBase::get_lmhead_full(cudaStream_t stream, NCCLCommunicator& comm, bool& accumulate) {
-    accumulate = !mIsFirstMicroStep;
-    return mFullNonBlock.LMHead;
-}
-
-Tensor& LLamaGradientsBlockShardedBase::get_lnf_w_full(cudaStream_t stream, NCCLCommunicator& comm, bool& accumulate) {
-    accumulate = !mIsFirstMicroStep;
-    return mFullNonBlock.LNF_w;
-}
-
-sLLamaBlockWeights<Tensor>& LLamaGradientsBlockShardedBase::get_block_full(int layer_idx, cudaStream_t stream, NCCLCommunicator& comm, bool& accumulate) {
-    accumulate = false;
-
-    auto& state = mGradStates.at(layer_idx % 2);
-    auto& dw = mGradBuffers.at(layer_idx % 2);
-    CUDA_CHECK(cudaStreamWaitEvent(stream, state.Event, 0)); // make sure the previous copy has finished
-    if (state.NeedsAccumulation) {
-        // already used; this means we need to schedule the accumulation first
-        sr_accumulate_layer(state.LayerIdx, mGradBuffers.at(state.LayerIdx % 2), mGradShards.at(state.LayerIdx), stream, comm);
-        state.NeedsAccumulation = false;
+    // if we have untied embeddings, we need to zero them out, same as above
+    if(mFullNonBlock.get_tensor(LLamaWeightID::LM_HEAD).Data != nullptr) {
+        fill_zero(mFullNonBlock.get_tensor(LLamaWeightID::EMBEDDING), stream);
     }
-    state.LayerIdx = layer_idx;
-    // reset local gradient buffers
-    fill_zero(dw.LN1_w, stream);
-    fill_zero(dw.LN2_w, stream);
-    fill_zero(dw.Attn_QKV_b, stream);
-    return dw;
-}
-
-void LLamaGradientsBlockShardedBase::notify_embeddings(cudaStream_t stream, NCCLCommunicator& comm) {
-    if(!mIsLastMicroStep) return;
-    scatter_reduce(mFullNonBlock.Embeddings, stream, mNonBlockEvent, comm);
-}
-
-void LLamaGradientsBlockShardedBase::notify_lmhead(cudaStream_t stream, NCCLCommunicator& comm) {
-    if(!mIsLastMicroStep) return;
-    if(mFullNonBlock.LMHead.Data == mFullNonBlock.Embeddings.Data) return;    // sync lmhead with embeddings
-    NvtxRange r{"notify_lmhead"};
-    scatter_reduce(mFullNonBlock.LMHead, stream, mNonBlockEvent, comm);
+    fill_zero(mFullNonBlock.get_tensor(LLamaWeightID::LNF_W), stream);
 }
 
-void LLamaGradientsBlockShardedBase::notify_lnf_w(cudaStream_t stream, NCCLCommunicator& comm) {
-    if(!mIsLastMicroStep) return;
-    scatter_reduce(mFullNonBlock.LNF_w, stream, mNonBlockEvent, comm);
+void LLamaGradientsBlockShardedBase::on_get_block(SimpleTensorContainer& block, cudaStream_t stream) {
+    fill_zero(block.get_tensor(LLamaWeightID::LN1_W), stream);
+    fill_zero(block.get_tensor(LLamaWeightID::LN2_W), stream);
+    fill_zero(block.get_tensor(LLamaWeightID::QKV_B), stream);
 }
 
-void LLamaGradientsBlockShardedBase::notify_block(int layer_idx, cudaStream_t stream, NCCLCommunicator& comm) {
-    auto& state = mGradStates[layer_idx % 2];
-    if(state.LayerIdx != layer_idx) {
-        throw std::logic_error("notify_block called with wrong layer index");
-    }
-    if (state.NeedsAccumulation) {
-        throw std::logic_error("notify_block called before accumulation has finished");
-    }
-
-    auto& dw = mGradBuffers.at(layer_idx % 2);
-    scatter_reduce(layer_idx, dw, stream, state.Event, comm);
-    state.NeedsAccumulation = true;
-}
-
-TensorShard& LLamaGradientsBlockShardedBase::get_embeddings_shard(cudaStream_t stream) {
-    return mNonBlockShards.Embeddings;
-}
-
-TensorShard& LLamaGradientsBlockShardedBase::get_lmhead_shard(cudaStream_t stream) {
-    return mNonBlockShards.LMHead;
-}
-
-TensorShard& LLamaGradientsBlockShardedBase::get_lnf_w_shard(cudaStream_t stream) {
-    return mNonBlockShards.LNF_w;
-}
-
-sLLamaBlockWeights<TensorShard>& LLamaGradientsBlockShardedBase::get_block_shard(int layer_idx, cudaStream_t stream) {
-    return mGradShards.at(layer_idx);
-}
 
 // ---------------------------------------------------------------------------------------------------------------------
 
-class LLamaGradientsBlockSharded_ScatterReduce : public LLamaGradientsBlockShardedBase {
-public:
-    using LLamaGradientsBlockShardedBase::LLamaGradientsBlockShardedBase;
-private:
-    void sr_accumulate_layer(int layer_idx,
-                             SimpleTensorContainer& dw,
-                             SimpleTensorContainer& sw,
-                             cudaStream_t stream,
-                             NCCLCommunicator& comm) override;
-};
+void LLamaGradientsBlockSharded_ScatterReduce::on_notify_block(int layer_idx, SimpleTensorContainer& block,
+    cudaStream_t stream, cudaEvent_t signal, NCCLCommunicator& comm) {
+    comm.reduce_scatter(block, stream, signal);
+}
 
 void LLamaGradientsBlockSharded_ScatterReduce::sr_accumulate_layer(int layer_idx,
                                                                    SimpleTensorContainer& dw,
@@ -349,7 +59,7 @@ void LLamaGradientsBlockSharded_ScatterReduce::sr_accumulate_layer(int layer_idx
                                                                    NCCLCommunicator& comm) {
     NvtxRange range("accumulate_layer", layer_idx);
     std::array<std::uint32_t, 8> rng;
-    mRng.generate(std::span(rng), 2*mStepCounter, layer_idx);
+    generate_rng(std::span(rng), 2*get_step_counter(), layer_idx);
 
     int rank = comm.rank();
     int world = comm.world_size();
@@ -357,7 +67,7 @@ void LLamaGradientsBlockSharded_ScatterReduce::sr_accumulate_layer(int layer_idx
     int i = 0;
     visit([&](Tensor& dst, Tensor& src){
         Tensor local_slice = shard_view(src, rank, world);
-        if(mIsFirstMicroStep) {
+        if(is_first_micro_step()) {
             CUDA_CHECK(cudaMemcpyAsync(dst.Data, local_slice.Data, local_slice.bytes(), cudaMemcpyDeviceToDevice, stream));
         } else {
             vector_add_sr(dst, dst, local_slice, 1.f, local_slice.nelem(), rng.at(i), stream);
@@ -368,19 +78,7 @@ void LLamaGradientsBlockSharded_ScatterReduce::sr_accumulate_layer(int layer_idx
 
 // ---------------------------------------------------------------------------------------------------------------------
 
-class LLamaGradientsBlockSharded_AllToAll : public LLamaGradientsBlockShardedBase {
-public:
-    using LLamaGradientsBlockShardedBase::LLamaGradientsBlockShardedBase;
-private:
-    void scatter_reduce(int layer_idx, SimpleTensorContainer& block, cudaStream_t stream, cudaEvent_t signal, NCCLCommunicator& comm) override;
-    void sr_accumulate_layer(int layer_idx,
-                             SimpleTensorContainer& dw,
-                             SimpleTensorContainer& sw,
-                             cudaStream_t stream,
-                             NCCLCommunicator& comm) override;
-};
-
-void LLamaGradientsBlockSharded_AllToAll::scatter_reduce(int layer_idx, SimpleTensorContainer& dw, cudaStream_t stream, cudaEvent_t signal, NCCLCommunicator& comm) {
+void LLamaGradientsBlockSharded_AllToAll::on_notify_block(int layer_idx, SimpleTensorContainer& dw, cudaStream_t stream, cudaEvent_t signal, NCCLCommunicator& comm) {
     auto& sw = get_block_shard(layer_idx, stream);
     int rank = comm.rank();
     int world = comm.world_size();
@@ -389,11 +87,11 @@ void LLamaGradientsBlockSharded_AllToAll::scatter_reduce(int layer_idx, SimpleTe
     {
         NvtxRange range("accumulate-own-shard", layer_idx);
         std::array<std::uint32_t, 8> rng;
-        mRng.generate(std::span(rng), 2*mStepCounter, layer_idx);
+        generate_rng(std::span(rng), 2*get_step_counter(), layer_idx);
         int i = 0;
         visit([&](Tensor& dst, Tensor& src){
             Tensor local_slice = shard_view(src, rank, world);
-            if(mIsFirstMicroStep) {
+            if(is_first_micro_step()) {
                 CUDA_CHECK(cudaMemcpyAsync(dst.Data, local_slice.Data, local_slice.bytes(), cudaMemcpyDeviceToDevice, stream));
             } else {
                 vector_add_sr(dst, dst, local_slice, 1.f, local_slice.nelem(), rng.at(i), stream);
@@ -403,7 +101,6 @@ void LLamaGradientsBlockSharded_AllToAll::scatter_reduce(int layer_idx, SimpleTe
     }
 
     // make sure we've done the local accumulation before we allow communication to begin.
-
     CUDA_CHECK(cudaEventRecord(signal, stream));
     NvtxRange range("all-to-all-gradients", layer_idx);
 
@@ -422,12 +119,12 @@ void LLamaGradientsBlockSharded_AllToAll::sr_accumulate_layer(int layer_idx,
     int rank = comm.rank();
     int world = comm.world_size();
     float scale = 1.f;
-    if (mIsLastMicroStep) {
+    if (is_last_micro_step()) {
         scale = 1.f / world;
     }
 
     std::array<std::uint32_t, 8> rng;
-    mRng.generate(std::span(rng), 2*mStepCounter, layer_idx);
+    generate_rng(std::span(rng), 2*get_step_counter(), layer_idx);
 
     int i = 0;
     visit([&](Tensor& s, Tensor& d){
@@ -436,18 +133,20 @@ void LLamaGradientsBlockSharded_AllToAll::sr_accumulate_layer(int layer_idx,
     }, sw, dw);
 }
 
-std::unique_ptr<LLamaGradsManager> LLamaGradsManager::create(std::uint64_t seed, int step, const TransformerConfig& config, const LLamaOptions& options, int rank, int world, const std::shared_ptr<TensorAllocator>& alloc) {
+std::unique_ptr<IGradientManager> create_grads_manager(std::uint64_t seed, int step, LLamaModel& model, const TransformerConfig& config,
+    const LLamaOptions& options, int rank, int world, const std::shared_ptr<TensorAllocator>& alloc)
+{
     if (options.ShardGradients) {
         if(options.UseAllToAllReduce) {
-            return std::make_unique<LLamaGradientsBlockSharded_AllToAll>(seed, step, config, options, rank, world, alloc);
+            return std::make_unique<LLamaGradientsBlockSharded_AllToAll>(config, model, seed, step, rank, world, options.OffloadGrads, alloc);
         } else {
-            return std::make_unique<LLamaGradientsBlockSharded_ScatterReduce>(seed, step, config, options, rank, world, alloc);
+            return std::make_unique<LLamaGradientsBlockSharded_ScatterReduce>(config, model, seed, step, rank, world, options.OffloadGrads, alloc);
         }
 
     } else {
         if(options.OffloadGrads) {
             throw std::logic_error("Offloading gradients is not supported for unsharded gradients");
         }
-        return std::make_unique<LLamaGradientsUnsharded>(seed, step, config, rank, world, alloc);
+        return std::make_unique<LLamaGradientsUnsharded>(config, model, seed, step, rank, world, alloc);
     }
 }
diff --git a/src/models/llama_gradients.h b/src/models/llama_gradients.h
index d677711..3b004b7 100644
--- a/src/models/llama_gradients.h
+++ b/src/models/llama_gradients.h
@@ -6,48 +6,54 @@
 #define LLMQ_SRC_MODELS_LLAMA_GRADIENTS_H
 
 #include "llama_weights.h"
-#include "utilities/philox.h"
+#include "training/gradients.h"
 
-class LLamaGradsManager {
+class LLamaModel;
+
+class LLamaGradientsUnsharded : public UnshardedGradientManager {
+public:
+    LLamaGradientsUnsharded(const TransformerConfig& cfg, IModel& model, std::uint64_t seed, int step, int rank,
+        int world, const std::shared_ptr<TensorAllocator>& alloc)
+        : UnshardedGradientManager(cfg, model, seed, step, rank, world, alloc) {
+    }
+
+    void on_first_micro_step(cudaStream_t stream) override;
+};
+
+class LLamaGradientsBlockShardedBase : public ShardedBlocksGradientManager {
+public:
+    using ShardedBlocksGradientManager::ShardedBlocksGradientManager;
+
+    void on_first_micro_step(cudaStream_t stream) override;
+    void on_get_block(SimpleTensorContainer& block, cudaStream_t stream) override;
+};
+
+class LLamaGradientsBlockSharded_ScatterReduce : public LLamaGradientsBlockShardedBase {
+public:
+    using LLamaGradientsBlockShardedBase::LLamaGradientsBlockShardedBase;
+private:
+    void on_notify_block(int layer_idx, SimpleTensorContainer& block, cudaStream_t stream, cudaEvent_t signal, NCCLCommunicator& comm) override;
+    void sr_accumulate_layer(int layer_idx,
+                             SimpleTensorContainer& dw,
+                             SimpleTensorContainer& sw,
+                             cudaStream_t stream,
+                             NCCLCommunicator& comm) override;
+};
+
+class LLamaGradientsBlockSharded_AllToAll : public LLamaGradientsBlockShardedBase {
 public:
-    virtual ~LLamaGradsManager() = default;
-
-    void start_micro_step(cudaStream_t stream, int micro_step, int total_steps);
-    virtual void end_micro_step(cudaStream_t stream, NCCLCommunicator& comm) = 0;
-
-    // Get references to full gradient accumulators for use in the backward pass
-    virtual Tensor& get_embeddings_full(cudaStream_t stream, NCCLCommunicator& comm, bool& accumulate) = 0;
-    virtual Tensor& get_lmhead_full(cudaStream_t stream, NCCLCommunicator& comm, bool& accumulate) = 0;
-    virtual Tensor& get_lnf_w_full(cudaStream_t stream, NCCLCommunicator& comm, bool& accumulate) = 0;
-    virtual sLLamaBlockWeights<Tensor>& get_block_full(int layer_idx, cudaStream_t stream, NCCLCommunicator& comm, bool& accumulate) = 0;
-
-    // Get references to sharded gradients for use in the optimizer
-    virtual TensorShard& get_embeddings_shard(cudaStream_t stream) = 0;
-    virtual TensorShard& get_lmhead_shard(cudaStream_t stream) = 0;
-    virtual TensorShard& get_lnf_w_shard(cudaStream_t stream) = 0;
-    virtual SimpleTensorContainer& get_block_shard(int layer_idx, cudaStream_t stream) = 0;
-
-    // notify that gradient calculations have been completed
-    virtual void notify_embeddings(cudaStream_t stream, NCCLCommunicator& comm) = 0;
-    virtual void notify_lmhead(cudaStream_t stream, NCCLCommunicator& comm) = 0;
-    virtual void notify_lnf_w(cudaStream_t stream, NCCLCommunicator& comm) = 0;
-    virtual void notify_block(int layer_idx, cudaStream_t stream, NCCLCommunicator& comm) = 0;
-
-    static std::unique_ptr<LLamaGradsManager> create(std::uint64_t seed, int step, const TransformerConfig& config,
-                                                     const LLamaOptions& options, int rank, int world,
-                                                     const std::shared_ptr<TensorAllocator>& alloc);
-
-protected:
-    LLamaGradsManager(std::uint64_t seed, int step);
-    virtual void on_first_micro_step(cudaStream_t stream) = 0;
-
-    void scatter_reduce(Tensor& tensor, cudaStream_t stream, cudaEvent_t signal, NCCLCommunicator& comm);
-    virtual void scatter_reduce(int layer_idx, SimpleTensorContainer& block, cudaStream_t stream, cudaEvent_t signal, NCCLCommunicator& comm);
-
-    Philox4x32 mRng;
-    int mStepCounter = -1;
-    bool mIsFirstMicroStep = true;
-    bool mIsLastMicroStep = false;
+    using LLamaGradientsBlockShardedBase::LLamaGradientsBlockShardedBase;
+private:
+    void on_notify_block(int layer_idx, SimpleTensorContainer& block, cudaStream_t stream, cudaEvent_t signal, NCCLCommunicator& comm) override;
+    void sr_accumulate_layer(int layer_idx,
+                             SimpleTensorContainer& dw,
+                             SimpleTensorContainer& sw,
+                             cudaStream_t stream,
+                             NCCLCommunicator& comm) override;
 };
 
+std::unique_ptr<IGradientManager> create_grads_manager(
+    std::uint64_t seed, int step, LLamaModel& model,
+    const TransformerConfig& config, const LLamaOptions& options,
+    int rank, int world, const std::shared_ptr<TensorAllocator>& alloc);
 #endif //LLMQ_SRC_MODELS_LLAMA_GRADIENTS_H
diff --git a/src/models/llama_model.cpp b/src/models/llama_model.cpp
index 2667333..c4e704d 100644
--- a/src/models/llama_model.cpp
+++ b/src/models/llama_model.cpp
@@ -389,7 +389,7 @@ void LLamaModel::backward(Tensor inputs, Tensor targets, NCCLCommunicator& comm,
     }
 
     bool accumulate;
-    auto& d_lnf_w = Grads->get_lnf_w_full(main_stream, comm, accumulate);
+    auto& d_lnf_w = Grads->get_non_block_full(LLamaWeightID::LNF_W, main_stream, comm, accumulate);
     Parameters->gather_lnf(comm);
     // backward the final layernorm
     rmsnorm_backward(rs->DActs[L-1].DResFFN.Value, d_lnf_w, rs->RMSNormScratch, rs->DActs[L - 1].DResFFN.Value, rs->DLNF,
@@ -398,7 +398,7 @@ void LLamaModel::backward(Tensor inputs, Tensor targets, NCCLCommunicator& comm,
     rs->release_res_ffn(L-1, main_stream);
 
     Parameters->release_lnf(main_stream);
-    Grads->notify_lnf_w(main_stream, comm);
+    Grads->notify_non_block(LLamaWeightID::LNF_W, main_stream, comm);
     rs->fetch_res_ffn(L-2, comm.stream());
     Parameters->gather_block(L - 1, comm, *rs);
     // now backward all the layers
@@ -429,22 +429,22 @@ void LLamaModel::backward(Tensor inputs, Tensor targets, NCCLCommunicator& comm,
 
         if(l > 0) {
             auto& prev_dacts = rs->DActs.at(l - 1);
-            rmsnorm_backward(prev_dacts.DResFFN.Value, dw.LN1_w, rs->RMSNormScratch, prev_dacts.DResAtt.Value, d_acts.DLN1,
+            rmsnorm_backward(prev_dacts.DResFFN.Value, dw.get_tensor(LLamaWeightID::LN1_W), rs->RMSNormScratch, prev_dacts.DResAtt.Value, d_acts.DLN1,
                              rs->get_res_ffn(l-1, main_stream), weights.LN1_w, rs->Acts[l].LN1_Rstd, prev_dacts.DResFFN.Quant.abs_max(),
                              B, T, C, rs->DeviceProp, main_stream);
             rs->release_res_ffn(l - 1, main_stream);
         } else {
-            rmsnorm_backward(rs->DEmb, dw.LN1_w, rs->RMSNormScratch, d_acts.DResAtt.Value, d_acts.DLN1,
+            rmsnorm_backward(rs->DEmb, dw.get_tensor(LLamaWeightID::LN1_W), rs->RMSNormScratch, d_acts.DResAtt.Value, d_acts.DLN1,
                              rs->Encoded, weights.LN1_w, rs->Acts[l].LN1_Rstd, nullptr, B, T, C, rs->DeviceProp, main_stream);
         }
         Parameters->release_block(l, main_stream);
         Grads->notify_block(l, main_stream, comm);
     }
 
-    auto& d_emb = Grads->get_embeddings_full(main_stream, comm, accumulate);
+    auto& d_emb = Grads->get_non_block_full(LLamaWeightID::EMBEDDING, main_stream, comm, accumulate);
     encoder_backward(d_emb, rs->EncoderBwdScratch, rs->EncoderBwdIndices, rs->EncoderBwdInfo,
                      rs->DEmb, rs->Inputs, inputs, B, T, C, OptimizerRNG(), main_stream, rs->SideStreamEvent, rs->SideStream);
-    Grads->notify_embeddings(main_stream, comm);
+    Grads->notify_non_block(LLamaWeightID::EMBEDDING, main_stream, comm);
 
     // make sure all gradients are communicated before we go to the update step.
     Grads->end_micro_step(main_stream, comm);
@@ -509,12 +509,21 @@ void LLamaModel::_backward_lmhead(long B, long T, float z_loss, int micro_step,
 
         // handle the LM-head. We run the d_lmhead matmul first, so that the gradient reduction can overlap with the DLNF matmul.
         bool accumulate;
-        auto& d_lmhead = Grads->get_lmhead_full(main_stream, comm, accumulate);
+        // get the correct matrix depending on whether we have tied embeddings
+        auto& d_lmhead = [&]() -> Tensor& {
+            if (Config.TiedWordEmbeddings) {
+                return Grads->get_non_block_full(LLamaWeightID::EMBEDDING, main_stream, comm, accumulate);
+            } else {
+                return Grads->get_non_block_full(LLamaWeightID::LM_HEAD, main_stream, comm, accumulate);
+            }
+        }();
+
+        // even if we overwrite for first micro-batch, we need to accumulate on non-first nano batch
         accumulate |= nano_step != 0;
         matmul(d_lmhead, lnf_slice, rs->Output, Tensor{}, nullptr, nullptr,
                rs->CublasLtHandle, rs->CuBlasWorkspace, C, V, nano_batch_size, EMMTranspose::NT, accumulate, main_stream, rs->MatmulBackend);
-        if (nano_step == nano_batches - 1) {
-            Grads->notify_lmhead(main_stream, comm);
+        if (nano_step == nano_batches - 1 && !Config.TiedWordEmbeddings) {
+            Grads->notify_non_block(LLamaWeightID::LM_HEAD, main_stream, comm);
         }
 
         matmul(dlnf_slice, Parameters->get_head(main_stream), rs->Output, Tensor{}, nullptr, nullptr,
@@ -610,8 +619,9 @@ void LLamaModel::_recompute_block(sLLamaBlockWeights<Tensor>& weights, sLLamaLay
     }
 }
 
-void LLamaModel::_backward_block(bool accumulate, sLLamaBlockWeights<Tensor>& weights, sLLamaGradBlock& d_weights,
+void LLamaModel::_backward_block(bool accumulate, sLLamaBlockWeights<Tensor>& weights, SimpleTensorContainer& d_weights,
                                  sLLamaLayerActivations& acts, sLLamaLayerGradients& d_acts) {
+    using namespace LLamaWeightID;
     auto& rs = RunState;
     cudaStream_t main_stream = rs->MainStream;
     // convenience shortcuts, size_t instead of int so that pointer arithmetics don't overflow
@@ -626,7 +636,7 @@ void LLamaModel::_backward_block(bool accumulate, sLLamaBlockWeights<Tensor>& we
     // backward the 2nd matmul of MLP
     // note that _recompute_block guarantees that if SwiGLu is already quantized (if necessary)
     rs->temp_acquire(d_acts.DSwiGLU);
-    backward_qmm(d_acts.DSwiGLU, d_weights.MLP_Down_w, Tensor{}, d_acts.DResFFN, acts.SwiGLu, weights.MLP_Down_w, Tensor{},
+    backward_qmm(d_acts.DSwiGLU, d_weights.get_tensor(DOWN_W), Tensor{}, d_acts.DResFFN, acts.SwiGLu, weights.MLP_Down_w, Tensor{},
                  accumulate, *rs, B, T, D, C, true, main_stream);
 
     swiglu_backward(d_acts.DMlpUp.Value, d_acts.DSwiGLU, acts.MlpUp, d_acts.DMlpUp.Quant.abs_max(), B, T, D, main_stream);
@@ -635,18 +645,18 @@ void LLamaModel::_backward_block(bool accumulate, sLLamaBlockWeights<Tensor>& we
     if(Options.grad_dtype() != d_acts.DMlpUp.Value.DType) {
         rs->temp_acquire(d_acts.DMlpUp.Quant);
     }
-    backward_qmm(d_acts.DLN2, d_weights.MLP_Up_w, Tensor{}, d_acts.DMlpUp, acts.LN2, weights.MLP_Up_w, Tensor{},
+    backward_qmm(d_acts.DLN2, d_weights.get_tensor(UP_W), Tensor{}, d_acts.DMlpUp, acts.LN2, weights.MLP_Up_w, Tensor{},
                  accumulate, *rs, B, T, C, 2 * D, !rs->Options.RecomputeRMSNorm, main_stream);
     if(Options.grad_dtype() != d_acts.DMlpUp.Value.DType) {
         rs->temp_free(d_acts.DMlpUp.Quant);
     }
 
     // rmsnorm backward does += to the dresidual, so it correctly accumulates grad from the MLP block above
-    rmsnorm_backward(d_acts.DResAtt.Value, d_weights.LN2_w, rs->RMSNormScratch, d_acts.DResFFN.Value, d_acts.DLN2,
+    rmsnorm_backward(d_acts.DResAtt.Value, d_weights.get_tensor(LN2_W), rs->RMSNormScratch, d_acts.DResFFN.Value, d_acts.DLN2,
                      acts.ResidualAtt, weights.LN2_w, acts.LN2_Rstd, d_acts.DResAtt.Quant.abs_max(), B, T, C, rs->DeviceProp, main_stream);
 
     bool recompute_ln1 = rs->Options.RecomputeRMSNorm || rs->Options.RecomputeAtt;
-    backward_qmm(d_acts.DAttY, d_weights.Attn_Out_w, Tensor{}, d_acts.DResAtt, acts.Att, weights.Attn_Out_w, Tensor{},
+    backward_qmm(d_acts.DAttY, d_weights.get_tensor(ATTO_W), Tensor{}, d_acts.DResAtt, acts.Att, weights.Attn_Out_w, Tensor{},
                  accumulate, *rs, B, T, C, C, false, main_stream);
 
     rs->temp_acquire(d_acts.DQKV.Value);
@@ -664,7 +674,7 @@ void LLamaModel::_backward_block(bool accumulate, sLLamaBlockWeights<Tensor>& we
     rs->temp_free(rs->CuDNNWorkspace);
     rope_backward(d_acts.DQKV.Value, d_acts.DQKV.Value, rs->FreqCis, d_acts.DQKV.Quant.abs_max(), B, T, Hq, Hkv, Hs, main_stream);
 
-    backward_qmm(d_acts.DLN1, d_weights.Attn_QKV_w, d_weights.Attn_QKV_b, d_acts.DQKV, acts.LN1, weights.Attn_QKV_w, rs->MatmulBiasScratch,
+    backward_qmm(d_acts.DLN1, d_weights.get_tensor(QKV_W), d_weights.get_tensor(QKV_B), d_acts.DQKV, acts.LN1, weights.Attn_QKV_w, rs->MatmulBiasScratch,
                  accumulate, *rs, B, T, C, Config.qkv_channels(), !recompute_ln1, main_stream);
     rs->temp_free(d_acts.DQKV.Value);
 }
@@ -756,6 +766,8 @@ void LLamaModel::fill_non_block_shapes(GenericTensorContainer& target, const Tra
     create(target.get_tensor(LLamaWeightID::LNF_W), C, 0, other_dtype);
     if(!config.TiedWordEmbeddings) {
         create(target.get_tensor(LLamaWeightID::LM_HEAD), V, C, matrix_dtype);
+    } else {
+        create(target.get_tensor(LLamaWeightID::LM_HEAD), 0, 0, matrix_dtype);
     }
 }
 
@@ -768,12 +780,12 @@ void LLamaModel::_calculate_gradient_norm(NCCLCommunicator& comm, float grad_cli
         global_norm_squared(rs->NormBuffer, grad, grad.nelem(), rs->DeviceProp, stream);
     };
 
-    norm_squared(Grads->get_embeddings_shard(stream));
+    norm_squared(Grads->get_non_block_shard(LLamaWeightID::EMBEDDING, stream));
 
     if(!Config.TiedWordEmbeddings) {
-        norm_squared(Grads->get_lmhead_shard(stream));
+        norm_squared(Grads->get_non_block_shard(LLamaWeightID::LM_HEAD, stream));
     }
-    norm_squared(Grads->get_lnf_w_shard(stream));
+    norm_squared(Grads->get_non_block_shard(LLamaWeightID::LNF_W, stream));
 
     for(int i = 0; i < Config.NumLayers; i++) {
         auto& block = Grads->get_block_shard(i, stream);
@@ -822,11 +834,11 @@ void LLamaModel::update(NCCLCommunicator& comm, float learning_rate, float beta_
     auto& nb_scales = OptimizerState->non_block_m_scales();
 
     using namespace LLamaWeightID;
-    run_update(Parameters->get_master_embeddings(), Grads->get_embeddings_shard(main_stream),
+    run_update(Parameters->get_master_embeddings(), Grads->get_non_block_shard(EMBEDDING, main_stream),
                OptimizerState->non_block_m().get_tensor(EMBEDDING), OptimizerState->non_block_v().get_tensor(EMBEDDING),
                nb_scales.get_tensor(EMBEDDING), weight_decay);
     comm.reduce_max(Parameters->get_master_embeddings().abs_max());
-    run_update(Parameters->get_master_lnf_w(), Grads->get_lnf_w_shard(main_stream),
+    run_update(Parameters->get_master_lnf_w(), Grads->get_non_block_shard(LNF_W, main_stream),
                OptimizerState->non_block_m().get_tensor(LNF_W), OptimizerState->non_block_v().get_tensor(LNF_W), nb_scales.get_tensor(LNF_W), 0.f);
     comm.reduce_max(Parameters->get_master_lnf_w().abs_max());
     CUDA_CHECK(cudaEventRecord(rs->OptEmbeddingsDone, main_stream));
@@ -868,7 +880,7 @@ void LLamaModel::update(NCCLCommunicator& comm, float learning_rate, float beta_
     }
 
     if(!Config.TiedWordEmbeddings) {
-        run_update(Parameters->get_master_lmhead(), Grads->get_lmhead_shard(main_stream),
+        run_update(Parameters->get_master_lmhead(), Grads->get_non_block_shard(LM_HEAD, main_stream),
                    OptimizerState->non_block_m().get_tensor(LM_HEAD), OptimizerState->non_block_v().get_tensor(LM_HEAD), nb_scales.get_tensor(LM_HEAD), weight_decay);
         comm.reduce_max(Parameters->get_master_lmhead().abs_max());
     }
@@ -914,7 +926,7 @@ void LLamaModel::allocate_run_state(const LLamaOptions& options, NCCLCommunicato
 
     {
         auto ctx = Allocator->with_context("Gradients");
-        Grads = LLamaGradsManager::create(42, 0, Config, options, comm.rank(), comm.world_size(), Allocator);
+        Grads = create_grads_manager(42, 0, *this, Config, options, comm.rank(), comm.world_size(), Allocator);
     }
 
     OptimizerRNG = std::minstd_rand{42};
diff --git a/src/models/llama_model.h b/src/models/llama_model.h
index 32c0bbe..d590c25 100644
--- a/src/models/llama_model.h
+++ b/src/models/llama_model.h
@@ -17,6 +17,8 @@
 
 // ---------------------------------------------------------------------------------------------------------------------
 
+class IGradientManager;
+
 struct LLamaOptions {
     bool KeepAllActivations = false;
     bool RecomputeSwiGLu = false;
@@ -113,7 +115,7 @@ class LLamaModel : public IModel {
     const TensorAllocator& get_allocator() const { return *Allocator; }
 
     const TransformerConfig& config() { return Config; }
-    LLamaGradsManager& grads() { return *Grads; }
+    IGradientManager& grads() { return *Grads; }
     LLamaRunState& run_state() { return *RunState; }
 
     IRunState& get_run_state() const override;
@@ -129,7 +131,7 @@ class LLamaModel : public IModel {
     void _forward_block(sLLamaBlockWeights<Tensor>& weights, sLLamaLayerActivations& activations, Tensor& residual);
     void _backward_lmhead(long B, long T, float z_loss, int micro_step, int grad_accum_steps, NCCLCommunicator& comm);
     void _recompute_block(sLLamaBlockWeights<Tensor>& weights, sLLamaLayerActivations& activations, Tensor& residual);
-    void _backward_block(bool accumulate, sLLamaBlockWeights<Tensor>& weights, sLLamaGradBlock& grads,
+    void _backward_block(bool accumulate, sLLamaBlockWeights<Tensor>& weights, SimpleTensorContainer& grads,
                          sLLamaLayerActivations& activations, sLLamaLayerGradients& d_activations);
 private:
     TransformerConfig Config;
@@ -137,7 +139,7 @@ class LLamaModel : public IModel {
     std::shared_ptr<TensorAllocator> Allocator;
     std::unique_ptr<LLamaWeightsManager> Parameters;
     std::unique_ptr<LLamaOptimizerStateManager> OptimizerState;
-    std::unique_ptr<LLamaGradsManager> Grads;
+    std::unique_ptr<IGradientManager> Grads;
     std::unique_ptr<LLamaRunState> RunState;
 
     std::minstd_rand OptimizerRNG;  //!< Seed generator for stochastic rounding in the optimizer
diff --git a/src/training/adamw_optimizer.cpp b/src/training/adamw_optimizer.cpp
index ae0b731..99889e7 100644
--- a/src/training/adamw_optimizer.cpp
+++ b/src/training/adamw_optimizer.cpp
@@ -12,28 +12,19 @@
 #include "utilities/stack.h"
 #include "utilities/lazy_allocator.h"
 
-static GenericTensorContainer shard_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) {
-        auto& t = c.get_tensor(i);
-        if (!t.empty()) { throw std::logic_error("shard_container called with non-empty tensor"); }
-        t.Sizes[0] = div_exact(t.Sizes[0], static_cast<long>(world));
-    }
-    return c;
-}
 
 AdamWStateManager::AdamWStateManager(TransformerConfig cfg, IModel& model, bool offload_m, bool offload_v,
     ETensorDType type_m, ETensorDType type_v, bool zero_copy, int rank, int world):
     mConfig(cfg), mOffloadM(offload_m), mOffloadV(offload_v), mUseZeroCopy(zero_copy), mRank(rank), mWorld(world), mMType(type_m), mVType(type_v) {
 
     if(mOffloadM && !mUseZeroCopy) {
-        mMDeviceBuffer[0] = shard_container(model.create_block_container(mConfig, mMType, mMType), mWorld);
-        mMDeviceBuffer[1] = shard_container(model.create_block_container(mConfig, mMType, mMType), mWorld);
+        mMDeviceBuffer[0] = shard_empty_container(model.create_block_container(mConfig, mMType, mMType), mWorld);
+        mMDeviceBuffer[1] = shard_empty_container(model.create_block_container(mConfig, mMType, mMType), mWorld);
     }
 
     if(mOffloadV && !mUseZeroCopy) {
-        mVDeviceBuffer[0] = shard_container(model.create_block_container(mConfig, mVType, mVType), mWorld);
-        mVDeviceBuffer[1] = shard_container(model.create_block_container(mConfig, mVType, mVType), mWorld);
+        mVDeviceBuffer[0] = shard_empty_container(model.create_block_container(mConfig, mVType, mVType), mWorld);
+        mVDeviceBuffer[1] = shard_empty_container(model.create_block_container(mConfig, mVType, mVType), mWorld);
     }
 
     if((mOffloadM || mOffloadV) && !mUseZeroCopy) {
@@ -173,11 +164,11 @@ void AdamWStateManager::allocate_state(IModel& model, cudaStream_t stream, EAllo
         LazyAllocator alloc_lazy;
         mBlocksM.resize(mConfig.NumLayers);
         for (int i = 0; i < mConfig.NumLayers; ++i) {
-            mBlocksM[i] = shard_container(model.create_block_container(mConfig, mMType, mMType), mWorld);
+            mBlocksM[i] = shard_empty_container(model.create_block_container(mConfig, mMType, mMType), mWorld);
             alloc_lazy.allocate(mBlocksM[i]);
             mStorageM.push_back(alloc_lazy.commit(alloc, mOffloadM ? kind : EAllocationType::ON_DEVICE, "m_block_shard"));
         }
-        mNonBlockM = shard_container(model.create_non_block_container(mConfig, mMType, mMType), mWorld);
+        mNonBlockM = shard_empty_container(model.create_non_block_container(mConfig, mMType, mMType), mWorld);
         alloc_lazy.allocate(mNonBlockM);
         mStorageM.push_back(alloc_lazy.commit(alloc, mOffloadM ? kind : EAllocationType::ON_DEVICE, "m_nonblock_shard"));
 
@@ -189,14 +180,14 @@ void AdamWStateManager::allocate_state(IModel& model, cudaStream_t stream, EAllo
         if(mMType == ETensorDType::FP8_E4M3) {
             // we "shard" for 128 as many GPUs, so that we get 1 scale per 128 weights.
             for (int i = 0; i < mConfig.NumLayers; ++i) {
-                mBlocksMScales[i] = shard_container(model.create_block_container(mConfig, ETensorDType::FP32, ETensorDType::FP32), 128 * mWorld);
+                mBlocksMScales[i] = shard_empty_container(model.create_block_container(mConfig, ETensorDType::FP32, ETensorDType::FP32), 128 * mWorld);
                 alloc_lazy.allocate(mBlocksMScales[i]);
                 alloc_lazy.commit(alloc, EAllocationType::ON_DEVICE, "m_block_scales");
                 visit([stream](Tensor& t){
                     fill_constant(t, 1.f, t.nelem(), stream);
                 }, mBlocksMScales[i]);
             }
-            mNonBlockMScales = shard_container(model.create_non_block_container(mConfig, ETensorDType::FP32, ETensorDType::FP32), 128 * mWorld);
+            mNonBlockMScales = shard_empty_container(model.create_non_block_container(mConfig, ETensorDType::FP32, ETensorDType::FP32), 128 * mWorld);
             alloc_lazy.allocate(mNonBlockMScales);
             alloc_lazy.commit(alloc, EAllocationType::ON_DEVICE, "m_nonblock_scales");
             visit([stream](Tensor& t){
@@ -215,11 +206,11 @@ void AdamWStateManager::allocate_state(IModel& model, cudaStream_t stream, EAllo
         LazyAllocator alloc_lazy;
         mBlocksV.resize(mConfig.NumLayers);
         for (int i = 0; i < mConfig.NumLayers; ++i) {
-            mBlocksV[i] = shard_container(model.create_block_container(mConfig, mVType, mVType), mWorld);
+            mBlocksV[i] = shard_empty_container(model.create_block_container(mConfig, mVType, mVType), mWorld);
             alloc_lazy.allocate(mBlocksV[i]);
             mStorageV.push_back(alloc_lazy.commit(alloc, mOffloadV ? kind : EAllocationType::ON_DEVICE, "v_block_shard"));
         }
-        mNonBlockV = shard_container(model.create_non_block_container(mConfig, mVType, mVType), mWorld);
+        mNonBlockV = shard_empty_container(model.create_non_block_container(mConfig, mVType, mVType), mWorld);
         alloc_lazy.allocate(mNonBlockV);
         mStorageV.push_back(alloc_lazy.commit(alloc, mOffloadV ? kind : EAllocationType::ON_DEVICE, "v_nonblock_shard"));
 
diff --git a/src/training/gradients.cpp b/src/training/gradients.cpp
new file mode 100644
index 0000000..55f5744
--- /dev/null
+++ b/src/training/gradients.cpp
@@ -0,0 +1,195 @@
+// Copyright (c) 2026, IST Austria, developed by Erik Schultheis
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "gradients.h"
+
+#include "model.h"
+#include "utilities/allocator.h"
+#include "utilities/comm.h"
+#include "utilities/lazy_allocator.h"
+
+IGradientManager::IGradientManager(std::uint64_t seed, int step) : mStepCounter(step), mRng(seed) {
+
+}
+
+void IGradientManager::start_micro_step(cudaStream_t stream, int micro_step, int total_steps) {
+    mIsFirstMicroStep = micro_step == 0;
+    mIsLastMicroStep = micro_step == total_steps - 1;
+    if (micro_step == 0) {
+        ++mStepCounter;
+        on_first_micro_step(stream);
+    }
+}
+
+UnshardedGradientManager::UnshardedGradientManager(const TransformerConfig& cfg, IModel& model, std::uint64_t seed, int step, int rank, int world, const std::shared_ptr<TensorAllocator>& alloc) : IGradientManager(seed, step){
+    LazyAllocator alloc_lazy;
+
+    mBlockGradients.resize(cfg.NumLayers);
+    mBlockShards.resize(cfg.NumLayers);
+    for (int i = 0; i < cfg.NumLayers; ++i) {
+        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(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(mNonBlockGradients, rank, world);
+
+    mGradEvent = create_named_event("grad_event");
+}
+
+UnshardedGradientManager::~UnshardedGradientManager() {
+    if (mGradEvent)
+        CUDA_CHECK(cudaEventDestroy(mGradEvent));
+}
+
+void UnshardedGradientManager::end_micro_step(cudaStream_t stream, NCCLCommunicator& comm) {
+    if (is_last_micro_step()) {
+        CUDA_CHECK(cudaStreamWaitEvent(stream, mGradEvent, 0));
+    }
+}
+
+Tensor& UnshardedGradientManager::get_non_block_full(std::size_t index, cudaStream_t stream, NCCLCommunicator& comm,
+    bool& accumulate) {
+    accumulate = !is_first_micro_step();
+    return mNonBlockGradients.get_tensor(index);
+}
+
+SimpleTensorContainer& UnshardedGradientManager::get_block_full(int layer_idx, cudaStream_t stream,
+    NCCLCommunicator& comm, bool& accumulate) {
+    accumulate = !is_first_micro_step();
+    return mBlockGradients.at(layer_idx);
+}
+
+Tensor& UnshardedGradientManager::get_non_block_shard(std::size_t index, cudaStream_t stream) {
+    return mNonBlockShards.get_tensor(index);
+}
+
+SimpleTensorContainer& UnshardedGradientManager::get_block_shard(int layer_idx, cudaStream_t stream) {
+    return mBlockShards.at(layer_idx);
+}
+
+void UnshardedGradientManager::notify_non_block(std::size_t index, cudaStream_t stream, NCCLCommunicator& comm) {
+    if(!is_last_micro_step()) return;
+    if (comm.world_size() != 1) {
+        NvtxRange r{"notify_non_block"};
+        comm.reduce_scatter(mNonBlockGradients.get_tensor(index), stream, mGradEvent);
+    }
+}
+
+void UnshardedGradientManager::notify_block(int layer_idx, cudaStream_t stream, NCCLCommunicator& comm) {
+    if(!is_last_micro_step()) return;
+    auto& dw = mBlockGradients.at(layer_idx);
+    if (comm.world_size() != 1) {
+        comm.reduce_scatter(dw, stream, mGradEvent);
+    }
+}
+
+ShardedBlocksGradientManager::ShardedBlocksGradientManager(const TransformerConfig& cfg, IModel& model, std::uint64_t seed, int step,
+    int rank, int world, bool offload, const std::shared_ptr<TensorAllocator>& alloc) : IGradientManager(seed, step)
+{
+    LazyAllocator alloc_lazy;
+    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(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);
+    alloc_lazy.allocate(mGradBuffers[0]);
+    alloc_lazy.allocate(mGradBuffers[1]);
+    alloc_lazy.commit(*alloc, EAllocationType::ON_DEVICE, "block_grad_buffers");
+
+    mGradStates[0].Event = create_named_event("grad_event_0");
+    mGradStates[1].Event = create_named_event("grad_event_1");
+    mNonBlockEvent = create_named_event("grad_nonblock_event");
+    mGradShards.resize(cfg.NumLayers);
+    for(int i = 0; i < cfg.NumLayers; ++i) {
+        EAllocationType kind = offload ? EAllocationType::PINNED : EAllocationType::ON_DEVICE;
+        mGradShards[i] = shard_view(model.create_block_container(cfg, cfg.DType, cfg.DType), rank, world);
+        alloc_lazy.allocate(mGradShards[i]);
+        alloc_lazy.commit(*alloc, kind, "block_grad_shards");
+    }
+}
+
+ShardedBlocksGradientManager::~ShardedBlocksGradientManager() {
+    if (mNonBlockEvent)
+        CUDA_CHECK(cudaEventDestroy(mNonBlockEvent));
+    for (auto& e : mGradStates) {
+        if (e.Event)
+            CUDA_CHECK(cudaEventDestroy(e.Event));
+    }
+}
+
+void ShardedBlocksGradientManager::end_micro_step(cudaStream_t stream, NCCLCommunicator& comm) {
+    for (int i = 0; i < 2; ++i) {
+        auto& state = mGradStates[i];
+        int layer_idx = state.LayerIdx;
+        if (state.NeedsAccumulation) {
+            // we need to wait for the previous accumulation to finish
+            CUDA_CHECK(cudaStreamWaitEvent(stream, state.Event, 0));
+            auto& dw = mGradBuffers.at(layer_idx % 2);
+            auto& sw = mGradShards.at(layer_idx);
+            sr_accumulate_layer(layer_idx, dw, sw, stream, comm);
+            state.NeedsAccumulation = false;
+        }
+    }
+    if (is_last_micro_step())
+        CUDA_CHECK(cudaStreamWaitEvent(stream, mNonBlockEvent, 0));
+}
+
+
+Tensor& ShardedBlocksGradientManager::get_non_block_full(std::size_t index, cudaStream_t stream, NCCLCommunicator& comm, bool& accumulate) {
+    accumulate = !is_first_micro_step();
+    return mFullNonBlock.get_tensor(index);
+}
+
+SimpleTensorContainer& ShardedBlocksGradientManager::get_block_full(int layer_idx, cudaStream_t stream, NCCLCommunicator& comm, bool& accumulate) {
+    accumulate = false;
+
+    auto& state = mGradStates.at(layer_idx % 2);
+    auto& dw = mGradBuffers.at(layer_idx % 2);
+    CUDA_CHECK(cudaStreamWaitEvent(stream, state.Event, 0)); // make sure the previous copy has finished
+    if (state.NeedsAccumulation) {
+        // already used; this means we need to schedule the accumulation first
+        sr_accumulate_layer(state.LayerIdx, mGradBuffers.at(state.LayerIdx % 2), mGradShards.at(state.LayerIdx), stream, comm);
+        state.NeedsAccumulation = false;
+    }
+    state.LayerIdx = layer_idx;
+    // reset local gradient buffers
+    on_get_block(dw, stream);
+    return dw;
+}
+
+void ShardedBlocksGradientManager::notify_non_block(std::size_t index, cudaStream_t stream, NCCLCommunicator& comm) {
+    if(!is_last_micro_step()) return;
+    NvtxRange r{"notify"};
+    comm.reduce_scatter(mFullNonBlock.get_tensor(index), stream, mNonBlockEvent);
+}
+
+void ShardedBlocksGradientManager::notify_block(int layer_idx, cudaStream_t stream, NCCLCommunicator& comm) {
+    auto& state = mGradStates[layer_idx % 2];
+    if(state.LayerIdx != layer_idx) {
+        throw std::logic_error("notify_block called with wrong layer index");
+    }
+    if (state.NeedsAccumulation) {
+        throw std::logic_error("notify_block called before accumulation has finished");
+    }
+
+    auto& dw = mGradBuffers.at(layer_idx % 2);
+    on_notify_block(layer_idx, dw, stream, state.Event, comm);
+    state.NeedsAccumulation = true;
+}
+
+
+Tensor& ShardedBlocksGradientManager::get_non_block_shard(std::size_t index, cudaStream_t stream) {
+    return mNonBlockShards.get_tensor(index);
+}
+
+SimpleTensorContainer& ShardedBlocksGradientManager::get_block_shard(int layer_idx, cudaStream_t stream) {
+    return mGradShards.at(layer_idx);
+}
diff --git a/src/training/gradients.h b/src/training/gradients.h
new file mode 100644
index 0000000..d5d36ed
--- /dev/null
+++ b/src/training/gradients.h
@@ -0,0 +1,183 @@
+// Copyright (c) 2026, IST Austria, developed by Erik Schultheis
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#ifndef LLMQ_TRAINING_GRADIENTS_H
+#define LLMQ_TRAINING_GRADIENTS_H
+
+#include <cstdint>
+#include <memory>
+#include <vector>
+
+#include "utilities/philox.h"
+#include "utilities/tensor_container.h"
+
+class IModel;
+class TensorAllocator;
+struct TransformerConfig;
+struct Tensor;
+class SimpleTensorContainer;
+class NCCLCommunicator;
+
+typedef struct CUstream_st *cudaStream_t;
+typedef struct CUevent_st* cudaEvent_t;
+
+/// Manages gradient accumulation and distribution across micro-batches and distributed workers.
+/// Subclasses support different gradient sharding strategies (e.g., unsharded, block-sharded) to trade off
+/// memory consumption and communication cost.
+///
+/// There are two sets of methods that allow querying gradient buffers. The `_full` methods return a view
+/// to an unsharded tensor to which the weight gradient of the current step needs to be added.
+/// The `_shard` methods return a view to a shard of the gradient tensor, e.g., for use in the optimizer.
+/// Shards may alias into the original full tensors, or they may use independent memory.
+/// These functions may block the supplied stream until the requested buffers are ready.
+class IGradientManager {
+public:
+    virtual ~IGradientManager() = default;
+
+    /// Must be called at the start of each micro-batch step.
+    void start_micro_step(cudaStream_t stream, int micro_step, int total_steps);
+
+    /// Must be called at the end of each micro-batch step. Finalizes gradient operations and may trigger
+    /// communication to reduce/scatter gradients across ranks.
+    virtual void end_micro_step(cudaStream_t stream, NCCLCommunicator& comm) = 0;
+
+    /// Get full (non-sharded) gradient tensor for non-block parameters (embeddings, LM head, etc.).
+    /// After the tensor has been updated, `notify_non_block` must be called to ensure proper synchronization.
+    /// \param index Weight index in the non-block parameter set
+    /// \param stream CUDA stream for potential synchronization
+    /// \param comm Communicator for potential gradient fetching
+    /// \param accumulate [out] true = add to existing gradients, false = overwrite
+    virtual Tensor& get_non_block_full(std::size_t index, cudaStream_t stream, NCCLCommunicator& comm, bool& accumulate) = 0;
+
+    /// Get full (non-sharded) gradient container for a transformer block.
+    /// After the tensors have been updated, `notify_block` must be called to ensure proper synchronization.
+    /// \param layer_idx Index of the transformer layer
+    /// \param stream CUDA stream for potential synchronization
+    /// \param comm Communicator for potential gradient fetching
+    /// \param accumulate [out] true = add to existing gradients, false = overwrite
+    virtual SimpleTensorContainer& get_block_full(int layer_idx, cudaStream_t stream, NCCLCommunicator& comm, bool& accumulate) = 0;
+
+    /// Get local rank's shard of non-block parameter gradients.
+    /// Used by optimizer to update the portion of parameters this rank is responsible for (ZeRO-1 style).
+    virtual Tensor& get_non_block_shard(std::size_t index, cudaStream_t stream) = 0;
+
+    /// Get local rank's shard of block parameter gradients.
+    /// Used by optimizer to update the portion of parameters this rank is responsible for (ZeRO-1 style).
+    virtual SimpleTensorContainer& get_block_shard(int layer_idx, cudaStream_t stream) = 0;
+
+    /// Notify that gradient computation is complete for a non-block parameter.
+    /// May trigger async communication (reduce-scatter/all-reduce) depending on sharding strategy.
+    virtual void notify_non_block(std::size_t index, cudaStream_t stream, NCCLCommunicator& comm) = 0;
+
+    /// Notify that gradient computation is complete for a transformer block.
+    /// May trigger async communication (reduce-scatter/all-reduce) depending on sharding strategy.
+    virtual void notify_block(int layer_idx, cudaStream_t stream, NCCLCommunicator& comm) = 0;
+
+    /// Gets the current step number
+    int get_step_counter() const { return mStepCounter; }
+
+protected:
+    IGradientManager(std::uint64_t seed, int step);
+
+    /// Check if this is the first micro-batch in the gradient accumulation loop.
+    /// Useful for determining whether to zero or accumulate gradients.
+    bool is_first_micro_step() const { return mIsFirstMicroStep; }
+
+    /// Check if this is the last micro-batch in the gradient accumulation loop.
+    /// Useful for determining when to finalize gradient communication.
+    bool is_last_micro_step() const { return mIsLastMicroStep; }
+
+    /// Generate deterministic random numbers, e.g., for stochastic rounding, using the philox generator.
+    template<std::size_t S>
+    void generate_rng(std::span<std::uint32_t, S> dst, std::uint32_t x, std::uint32_t y) {
+        mRng.generate(dst, x, y);
+    }
+
+private:
+    /// Hook called on the first micro-batch step for gradient initialization (e.g., zeroing).
+    virtual void on_first_micro_step(cudaStream_t stream) = 0;
+
+    int mStepCounter = -1;
+    bool mIsFirstMicroStep = true;
+    bool mIsLastMicroStep = false;
+    Philox4x32 mRng;
+};
+
+/// Gradient manager where each worker maintains full gradient copies.
+/// Implements ZeRO-1 optimization: gradients are replicated across all ranks,
+/// but each rank's optimizer is responsible for updating only its assigned shard.
+/// This reduces optimizer memory while keeping gradient computation simple.
+class UnshardedGradientManager : public IGradientManager {
+public:
+    ~UnshardedGradientManager();
+    void end_micro_step(cudaStream_t stream, NCCLCommunicator& comm) override;
+
+    Tensor& get_non_block_full(std::size_t index, cudaStream_t stream, NCCLCommunicator& comm, bool& accumulate) override;
+    SimpleTensorContainer& get_block_full(int layer_idx, cudaStream_t stream, NCCLCommunicator& comm, bool& accumulate) override;
+
+    Tensor& get_non_block_shard(std::size_t index, cudaStream_t stream) override;
+    SimpleTensorContainer& get_block_shard(int layer_idx, cudaStream_t stream) override;
+
+    void notify_non_block(std::size_t index, cudaStream_t stream, NCCLCommunicator& comm) override;
+    void notify_block(int layer_idx, cudaStream_t stream, NCCLCommunicator& comm) override;
+
+protected:
+    UnshardedGradientManager(const TransformerConfig& cfg, IModel& model, std::uint64_t seed, int step, int rank, int world, const std::shared_ptr<TensorAllocator>& alloc);
+
+    cudaEvent_t mGradEvent = nullptr;  ///< Synchronization event for backward pass completion
+
+    /// Full gradient accumulators for transformer blocks (replicated across all ranks)
+    std::vector<GenericTensorContainer> mBlockGradients;
+    /// Full gradient accumulators for non-block parameters (replicated across all ranks)
+    GenericTensorContainer mNonBlockGradients;
+
+    /// Views into mBlockGradients representing this rank's shard
+    std::vector<GenericTensorContainer> mBlockShards;
+    /// Views into mNonBlockGradients representing this rank's shard
+    GenericTensorContainer mNonBlockShards;
+};
+
+/// Gradient manager where the transformer blocks are sharded across
+/// workers, but the non-block weights are replicated.
+class ShardedBlocksGradientManager : public IGradientManager {
+public:
+    ShardedBlocksGradientManager(const TransformerConfig& cfg, IModel& model, std::uint64_t seed, int step, int rank, int world, bool offload, const std::shared_ptr<TensorAllocator>& alloc);
+    ~ShardedBlocksGradientManager();
+
+    void end_micro_step(cudaStream_t stream, NCCLCommunicator& comm) override;
+
+    Tensor& get_non_block_full(std::size_t index, cudaStream_t stream, NCCLCommunicator& comm, bool& accumulate) override;
+    SimpleTensorContainer& get_block_full(int layer_idx, cudaStream_t stream, NCCLCommunicator& comm, bool& accumulate) override;
+
+    Tensor& get_non_block_shard(std::size_t index, cudaStream_t stream) override;
+    SimpleTensorContainer& get_block_shard(int layer_idx, cudaStream_t stream) override;
+
+    void notify_non_block(std::size_t index, cudaStream_t stream, NCCLCommunicator& comm) override;
+    void notify_block(int layer_idx, cudaStream_t stream, NCCLCommunicator& comm) override;
+protected:
+    virtual void sr_accumulate_layer(int layer_idx,
+                                     SimpleTensorContainer& dw,
+                                     SimpleTensorContainer& sw,
+                                     cudaStream_t stream,
+                                     NCCLCommunicator& comm) = 0;
+
+    virtual void on_get_block(SimpleTensorContainer& block, cudaStream_t stream) = 0;
+    virtual void on_notify_block(int layer_idx, SimpleTensorContainer& block, cudaStream_t stream, cudaEvent_t signal, NCCLCommunicator& comm) = 0;
+
+    GenericTensorContainer mFullNonBlock;
+    GenericTensorContainer mNonBlockShards;
+
+    std::array<GenericTensorContainer, 2> mGradBuffers;
+    struct sBlockState {
+        cudaEvent_t Event;
+        int LayerIdx = -1;
+        bool NeedsAccumulation = false;
+    };
+    std::array<sBlockState, 2> mGradStates;
+    std::vector<GenericTensorContainer> mGradShards;
+    cudaEvent_t mNonBlockEvent;
+};
+
+
+#endif //LLMQ_TRAINING_GRADIENTS_H
diff --git a/src/utilities/comm.cpp b/src/utilities/comm.cpp
index 7d5c96f..cebaa6e 100644
--- a/src/utilities/comm.cpp
+++ b/src/utilities/comm.cpp
@@ -15,6 +15,7 @@
 #include "gpu_info.h"
 #include "kernels/kernels.h"
 #include "tensor.h"
+#include "tensor_container.h"
 #include "utils.h"
 
 void nccl_check(ncclResult_t status, const char* file, int line) {
@@ -164,6 +165,18 @@ void NCCLCommunicator::execute_transaction(cudaEvent_t signal) {
     mCmdBuf->Commands.clear();
 }
 
+void NCCLCommunicator::reduce_scatter(Tensor& tensor, cudaStream_t stream, cudaEvent_t signal) {
+    begin_transaction(stream);
+    schedule_reduce_scatter(tensor);
+    execute_transaction(signal);
+}
+
+void NCCLCommunicator::reduce_scatter(SimpleTensorContainer& container, cudaStream_t stream, cudaEvent_t signal) {
+    begin_transaction(stream);
+    visit([&](Tensor& t){ schedule_reduce_scatter(t); }, container);
+    execute_transaction(signal);
+}
+
 void NCCLCommunicator::schedule_reduce_scatter(Tensor& tensor) {
     if(tensor.empty()) return;
     mCmdBuf->Commands.emplace_back(CommandBuffer::ScatterReduce{.DType = tensor.DType, .Tensor = tensor.Data, .Elements = tensor.nelem()});
diff --git a/src/utilities/comm.h b/src/utilities/comm.h
index 7e61f59..e293dbc 100644
--- a/src/utilities/comm.h
+++ b/src/utilities/comm.h
@@ -13,6 +13,8 @@
 
 #include <cuda_bf16.h>
 
+class SimpleTensorContainer;
+
 namespace std
 {
     class jthread;
@@ -50,6 +52,9 @@ class NCCLCommunicator {
     void schedule_destructive_all_to_all(Tensor& tensor);
     void execute_transaction(cudaEvent_t signal);
 
+    void reduce_scatter(Tensor& tensor, cudaStream_t stream, cudaEvent_t signal);
+    void reduce_scatter(SimpleTensorContainer& container, cudaStream_t stream, cudaEvent_t signal);
+
     void reduce_mean(float* values, int n, cudaStream_t stream);
     void reduce_norm(float* norm_squared, cudaStream_t stream);
 
diff --git a/src/utilities/tensor.cpp b/src/utilities/tensor.cpp
index 2fa0107..f434b68 100644
--- a/src/utilities/tensor.cpp
+++ b/src/utilities/tensor.cpp
@@ -157,3 +157,21 @@ std::size_t GenericTensorContainer::num_tensors() const noexcept {
 const Tensor& GenericTensorContainer::get_tensor(std::size_t idx) const {
     return mTensors.at(idx);
 }
+
+GenericTensorContainer shard_empty_container(GenericTensorContainer&& c, int world) {
+    // can't use visit here, because we explicitly want to iterate over empty tensors
+    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_empty_container called with non-empty tensor"); }
+        t.Sizes[0] = div_exact(t.Sizes[0], static_cast<long>(world));
+    }
+    return std::move(c);
+}
+
+GenericTensorContainer shard_view(const GenericTensorContainer& c, int rank, int world) {
+    std::vector<Tensor> shards(c.num_tensors());
+    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};
+}
diff --git a/src/utilities/tensor_container.h b/src/utilities/tensor_container.h
index 7d5924e..fa89b88 100644
--- a/src/utilities/tensor_container.h
+++ b/src/utilities/tensor_container.h
@@ -57,6 +57,13 @@ class GenericTensorContainer final : public SimpleTensorContainer {
     std::vector<Tensor> mTensors;
 };
 
+//! 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);
+
+//! 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);
 
 class ITensorContainer {
   public: