Method run_step

csrc/adam/cpu_adam_impl.cpp:572–609

private:
    void run_step(int now_state, int64_t step, const std::vector<ZenHP>& hps)
    {
        auto opt = std::static_pointer_cast<Adam_Optimizer>(s_optimizers[opt_id_]);
        for (size_t g = 0; g < groups_.size(); ++g) {
            const ZenHP& hp = hps[g];
            // Groups share one Adam_Optimizer; advance its bias-correction state for
            // this group before the pool reads it (pool is idle here -> no race).
            opt->IncrementStep(step, hp.beta1, hp.beta2);
            opt->update_state(hp.lr, hp.eps, hp.weight_decay, hp.bias_correction);

            ZenGroup& grp = groups_[g];
            torch::Tensor& P = grp.param;
            torch::Tensor& G = grp.grad[now_state];
            torch::Tensor& M = grp.exp_avg[now_state];
            torch::Tensor& V = grp.exp_avg_sq[now_state];

            auto it = invokers.find(std::tuple(P.scalar_type(), M.scalar_type()));
            TORCH_CHECK(it != invokers.end(),
                        "ZenFlowAdam: unsupported param/state dtype combination");
            auto fn = it->second;

            char* pp = static_cast<char*>(P.data_ptr());
            char* gp = static_cast<char*>(G.data_ptr());
            char* mp = static_cast<char*>(M.data_ptr());
            char* vp = static_cast<char*>(V.data_ptr());
            char* sp = grp.stale.defined() ? static_cast<char*>(grp.stale.data_ptr()) : nullptr;
            const size_t pe = P.element_size();
            const size_t se = M.element_size();
            const size_t numel = P.numel();

            pool_->parallel_for(numel, kZenAdamAlign, [=](size_t b, size_t e) {
                const size_t len = e - b;
                // parallel=false: each pinned thread runs its slice serially.
                fn(opt, pp + b * pe, gp + b * pe, mp + b * se, vp + b * se, len, false);
                if (sp) std::memcpy(sp + b * pe, pp + b * pe, len * pe);
            });
        }
    }

    int opt_id_;
    std::vector<ZenGroup> groups_;