Diff Coverage

                real_id = self._ctx.real_lay_ids[chunk_id][i][lay_id]
                name = f"Lay_{real_id}_{lay_type.name[0]}"
                idx = chunk_id * len(dyn_mem_i) + lay_id
                stat = self.mb(stat_mem_i[chunk_id][lay_id])
                dyn = self.mb(sum(dyn_mem_i[chunk_id][lay_id]))
                ax.bar(
                    0.1,
                    [dyn],
                    bottom=bottoms["dyn"],

                    sm["stat"][stage_id][chunk_id][lay_id] = static_mem
                    dyn_val = self._inner_dynamic_mem()
                    sm["dyn"][stage_id][chunk_id][lay_id] = dyn_val
                    if os.environ.get("_SAPP_DEBUG_DYN") == "1":
                        _sys.stderr.write(
                            f"DEBUG_DYN stage={stage_id} chunk={chunk_id} lay={lay_id} "
                            f"node={node} active={dyn_val[0] if isinstance(dyn_val, tuple) else dyn_val} "
                            f"comm={dyn_val[1] if isinstance(dyn_val, tuple) else 0} "
                            f"mf={self._ctx.micro_factor} h={self._ccfg.h} "

        if fw_oh > 0 and d_shard > 1:
            if d_replicate > 1:
                fw_oh = int(fw_oh * (1.0 + 0.03 * (d_shard - 1)))
            else:
                fw_oh = int(fw_oh * (1.0 + 0.01 * (d_shard - 1)))
        if fw_oh > 0 and d_replicate > 1:
            fw_oh = int(fw_oh / (1.0 + 0.025 * (d_replicate - 1)))
        if fw_oh > 0 and d_shard > 1 and d_replicate > 1:
            fw_oh = int(fw_oh / (1.0 + 0.03 * (d_shard - 1) * (d_replicate - 1) / d_shard))

            fw_oh = int(fw_oh / (1.0 + 0.025 * (d_replicate - 1)))
        if fw_oh > 0 and d_shard > 1 and d_replicate > 1:
            fw_oh = int(fw_oh / (1.0 + 0.03 * (d_shard - 1) * (d_replicate - 1) / d_shard))
        if fw_oh > 0 and d_replicate > 1 and tp >= 8:
            tp_hsdp_extra = int(fw_oh * 0.28)
            fw_oh += tp_hsdp_extra
        return fw_oh

    def __build_stage_insight(
        self, stage_id, raw_dyn, total_raw, total_fw_oh,

            fw_oh_share = total_fw_oh * (
                0.5 / num_stages + 0.5 * raw_dyn[stage_id] / total_raw
            )
            ins["Dynamic"] += int(fw_oh_share)
        elif ins["Dynamic"] > 0:
            ins["Dynamic"] += total_fw_oh // num_stages
        stage_accu = sm["logs"][stage_id].accu_mem_type
        ins["ModelParameters"] = self.mb(stage_accu[MemType.MODEL_PARAM])
        ins["OptimizerStates"] = self.mb(stage_accu[MemType.OPTIM_STATE])
        ins["AccumulGradients"] = self.mb(stage_accu[MemType.ACCU_GRAD])

        ins["AllGather Comm"] = self.mb(stage_accu[MemType.AG_COMM])
        ins["All2All Comm"] = self.mb(stage_accu[MemType.A2A_COMM])

        if self._ccfg.cp > 1:
            cp_memory = EvalBody.act_cp_layer(self._ccfg, self._ctx)
            cp_comm_buffer = EvalLayerComm.cp_comm_buffer(self._ccfg, self._ctx)

            ins["CP KV Cache"] = self.mb(cp_memory.kv_cache_memory)
            ins["CP Attn Scores"] = self.mb(cp_memory.attention_scores_memory)
            ins["CP Softmax"] = self.mb(cp_memory.softmax_outputs_memory)
            ins["CP Comm Buffer"] = self.mb(cp_comm_buffer)
            ins["CP Reduction"] = self.mb(cp_memory.total_reduction)

        ins["Node Log"] = sm["logs"][stage_id].node_compute_log
        if verbose and spec_stage_id in (-1, stage_id):
            self.__verbose_insights(sm, stage_id, ins)
        ins["Static"] = self.mb(ins["Static"])
        ins["Dynamic"] = self.mb(ins["Dynamic"])
        insights.append(ins)

        num_stages = self._ccfg.p
        safety_buffer = 1024 * 1024 * 1024  # 1 GB
        raw_dyn = self.__compute_raw_dynamic(num_stages, stages, record_lay_types, sm, safety_buffer)
        if os.environ.get("_SAPP_DEBUG_DYN") == "1":
            for sid in range(num_stages):
                act_s = sum(
                    sum(aval for aval, _ in cell) for cell in sm["dyn"][sid]
                )
                comm_p = max(
                    max(comm for _, comm in cell) for cell in sm["dyn"][sid]
                )
                _sys.stderr.write(
                    f"DEBUG_POST stage={sid} act_sum={act_s/(1024**2):.1f}MB "
                    f"comm_peak={comm_p/(1024**2):.1f}MB "
                    f"raw_dyn={raw_dyn[sid]/(1024**2):.1f}MB\n"
                )

        current_node = getattr(self._ctx, "current_node", None)
        if current_node == LayerType.OUTPUT_LAYER:
            t = getattr(self._ccfg, "t", 1)
            if t > 1:
                bwd_workspace = (
                    (1.0 - 1.0 / t)
                    * getattr(self._ccfg, "s", 0)
                    * getattr(self._ccfg, "b", 0)
                    * getattr(self._ccfg, "v", 0)

                        self._ctx.current_lay_id = (
                            f"G_{self._ctx.current_lay_id}"
                        )
                        self._ctx.current_node = last_mtp
                        res += self.__bwd_dyn_mem(
                            is_recomputed=False, default_micro_factor=1
                        )
        return res

        Result is in bytes (like TP activation comm), unlike CP/DP which are
        in element counts (parameter comm).
        """
        if ccfg.comm_ep == 0:
            return 0
        if ccfg.n_exp <= 1:
            return 0
        if ccfg.ep <= 1:
            return 0

            Peak TP all-gather buffer size in bytes.
        """
        if ccfg.t <= 1 or ccfg.comm_t == 0:
            return 0
        rec_layer = ctx.current_node == LayerType.SEL_REC_LAYER
        rec_coeff = EvalUtils.rec_coeff(rec_layer, ccfg.rec_op.gather)

        attn_buf = ccfg.s * ccfg.b * ccfg.h * ccfg.bytes_compute
        if ccfg.n_exp > 1:
            ffn_buf = ccfg.s * ccfg.b * 2 * ccfg.hff * ccfg.bytes_compute
        else:
            ffn_buf = ccfg.s * ccfg.b * 2 * ccfg.hff * ccfg.bytes_compute

        peak_buf = max(attn_buf, ffn_buf)
        return rec_coeff * ccfg.comm_t * mb * peak_buf / ccfg.cp

        embedding gradient fragments.
        """
        if getattr(ccfg, "comm_hsdp", 0) <= 0:
            return 0.0
        param_size = ctx.eval.num_p(ccfg, ctx)
        d_shard = ccfg.d_shard_or_d
        sharded_size = param_size / (d_shard * ccfg.cp * ccfg.t)
        return ccfg.comm_hsdp * sharded_size * ccfg.bytes_compute

    @staticmethod
    def fsdp_grad_comm_embed(ccfg: CostModelConfig, ctx: Context) -> float:
        """FSDP/HSDP gradient reduce-scatter buffer for embedding params (bytes).

    def __eq__(self, other: object) -> bool:
        """Check equality based on pp and vpp values."""
        if not isinstance(other, Pipeline):
            return NotImplemented
        return self.pp == other.pp and self.vpp == other.vpp


def infer_pp_and_vpp(offset: Any) -> Any:

            self.__maybe_set_int(target_ccfg, attr, value)
        if target_ccfg.use_seq_parallel:
            target_ccfg.sp = target_ccfg.t if target_ccfg.t > 1 else 1
        else:
            target_ccfg.sp = 1
        if target_ccfg.recompute_slice_activation and tp is not None:
            target_ccfg.shard_recompute_input = target_ccfg.t
            target_ccfg.shard_output_activ = target_ccfg.t
        if op and isinstance(op, int):

        if fsdp_kw is True:
            target_ccfg.fsdp = True
            target_ccfg.has_op = True
        elif fsdp_kw is False:
            target_ccfg.fsdp = False
            if op and isinstance(op, int) and op <= 1:
                target_ccfg.has_op = False
            elif not op:
                target_ccfg.has_op = False

    def __apply_d_shard_strategy(self, target_ccfg, d_shard_kw):
        """Apply the HSDP d_shard value, defaulting from dp/fsdp.

            and isinstance(d_shard_kw, int)
            and 1 < d_shard_kw < target_ccfg.d
        )
        if is_hsdp:
            target_ccfg.fsdp = True
        saved_d_shard = target_ccfg.d_shard

        logger.debug(
            "in ccfg: DP = %d, TP(MP) = %d, EP = %d, CP = %d, "

            target_ccfg.fsdp,
        )
        self.__apply_parser_config(target_ccfg)
        if is_hsdp:
            target_ccfg.d_shard = saved_d_shard
        self.__apply_recompute_offset(target_ccfg, off, fr, sr)
        self.__maybe_set_int(target_ccfg, "cp", cp)

    def get_strategy(self) -> Dict[str, Any]:

        self.ccfg.bytes_p = self.ccfg.fp_bytes(
            self.config.model.model_config.param_init_type,
        )
        if not self.ccfg.bytes_p:
            self.ccfg.bytes_p = self.ccfg.fp_bytes(
                self.config.model.model_config.params_dtype
            )
        self.ccfg.bytes_compute = self.ccfg.fp_bytes(
            self.config.model.model_config.compute_dtype

        self.ccfg.bytes_softmax = self.ccfg.fp_bytes(
            self.config.model.model_config.softmax_compute_type
        )
        if not self.ccfg.bytes_softmax:
            self.ccfg.bytes_softmax = self.ccfg.fp_bytes(
                self.config.model.model_config.softmax_compute_dtype
            )
        if not self.ccfg.bytes_p:
            raise AttributeError("bytes_p not positive")
        if not self.ccfg.bytes_compute:
            raise AttributeError("bytes_compute not positive")
        self.ccfg.bytes_grad = 4
        self.ccfg.bytes_os = 4
        self.ccfg.bytes_norm = 4

            self.ccfg.framework_overhead = int(
                10.5 * 1024 ** 3 * (cap_gb / 32.0) ** 1.74
            )
        else:
            self.ccfg.framework_overhead = getattr(
                self.ccfg, "framework_overhead", 0
            )

    def __config_parse_yaml_optimizer_shard(self):

            name = "EP"
        elif self == self.CP_COMM:
            name = "CP"
        elif self == self.FSDP_COMM:
            name = "FSDP"
        elif self == self.PP_COMM:
            name = "P2P"
        elif self == self.BUBBLE:
            name = "BBL"

        parts[RealParts.DP_WAIT][-1] += real["op_wait"]

    tp = "tp_wait"
    if tp in real.keys():
        parts[RealParts.MP_WAIT][-1] += real["tp_wait"]

    sp = "sp_wait"
    if sp in real.keys():
        parts[RealParts.MP_WAIT][-1] += real["sp_wait"]

    Returns:
        bool: The parsed boolean.
    """
    if isinstance(value, str):
        return value.strip().lower() not in ("false", "0", "no", "off", "")
    return bool(value)


class Dimension:
    """Output dimension"""

            valid = dims_val[MBN] >= dims_val[PP]
            valid = valid and not (dims_val[PP] == 1 and dims_val[MBN] > 1)
            if not valid:
                return False
        return True

    @staticmethod
    def _check_power_of_two(dim, value):
        """Check whether *value* is a power of 2 for the given dimension.

        """Validate that OP is a power of 2."""
        if OP in self.all_dims:
            op = self.dims_val[OP]
            if not (op & (op - 1)) == 0:
                logger.warning("OP %d must be a power of 2", op)
                return False
        return True

    def _validate_fsdp_constraints(self):
        """Validate FSDP coexistence with OP, EP and the DP lower bound."""

        """Validate FSDP coexistence with OP, EP and the DP lower bound."""
        if not (self.has_dim(FSDP) and self.dims_val[FSDP]):
            return True
        if self.has_dim(OP) and self.dims_val[OP] > 1:
            logger.warning("FSDP and OP cannot coexist (OP > 1)")
            return False
        if self.has_dim(EP) and self.dims_val[EP] > 1:
            logger.warning("FSDP and EP cannot coexist")
            return False
        if self.has_dim(DP) and self.dims_val[DP] < 2:
            logger.warning("FSDP requires DP >= 2")
            return False
        return True

    def _validate_hsdp_fsdp_coexistence(self, d_shard):
        """Validate HSDP constraints that depend on FSDP being present."""

    def _validate_hsdp_fsdp_coexistence(self, d_shard):
        """Validate HSDP constraints that depend on FSDP being present."""
        if d_shard > 1 and not self.dims_val[FSDP]:
            logger.warning("HSDP d_shard > 1 requires FSDP enabled")
            return False
        if d_shard > 1 and self.dims_val[EP] > 1:
            logger.warning("HSDP and EP cannot coexist")
            return False
        if d_shard > 1 and self.dims_val[OP] > 1:
            logger.warning("HSDP and OP cannot coexist (OP > 1)")
            return False
        return True

    def _validate_hsdp_constraints(self):
        """Validate HSDP-related dimension constraints."""

            return False
        if self.has_dim(FSDP) and not self._validate_hsdp_fsdp_coexistence(
            d_shard
        ):
            return False
        if d_shard > 1 and not (d_shard & (d_shard - 1)) == 0:
            logger.warning("HSDP d_shard %d must be a power of 2", d_shard)
            return False
        return True

    def is_valid(self) -> bool:

        d_shard=1 is excluded because it means "no sharding" which
        contradicts fsdp=True.  When this path is not taken (fsdp=False
        or HSDP dimension absent), d_shard defaults to 1 elsewhere.
        """
        if Dim.HSDP not in self.dimensions:
            return [dp]
        power_of_2_divisors = Hard.all_divisors(dp)
        valid = [s for s in power_of_2_divisors if 2 <= s < dp]
        if not valid:
            return [dp]
        return [dp] + valid

    def max_op(self, dp: int, tp: int, ep: int) -> int:
        """Compute bound for dimension OP"""
        if (

    def _parallel_loops_fsdop(self, space, pool, dtpc_p, mbsn, ep, vpp, dp, tp):
        """Iterate over FSDP/OP/d_shard/SP sub-space for one (ep, vpp)."""
        for fsdp in self.config.bool_space(Dim.FSDP):
            if fsdp and dp < 2:
                continue
            if fsdp and ep > 1:
                continue
            if fsdp:
                op_space = [1]
                d_shard_space = self.config.hsdp_space(dp)
            else:
                op_space = self.config.space(
                    Dim.OP, self.config.max_op(dp, tp, ep)
                )

                        ),
                    )
                else:
                    if self.enable_debug:
                        debugger = Debug.Debug(
                            config,
                            info_type=Debug.PerfParts,
                            enable=self.enable_debug,
                        )
                        score = estimate_performance(
                            self.config.ccfg,
                            debugger=debugger,
                            device_type=self.machine.device,
                            memory=mem,

                            device_type=self.machine.device,
                            memory=mem,
                            cache_file=cache_file,
                        )
                        debugger.write()
                        debugger.info.pop(Debug.PerfParts.TOTAL, None)
                        debugger.info.pop(Debug.PerfParts.MEMORY, None)
                        debug_parts = list(debugger.info.keys())
                        values = list(debugger.info.values())
                    else:
                        score = estimate_performance(
                            self.config.ccfg,
                            device_type=self.machine.device,

        # If MB (microbatch number) is not in the config dimensions,
        # inject it so set_strategy can compute m from gbs/d/b.
        injected_mbn = False
        if Dim.MBN not in config.dims_val and Dim.MBS in config.dims_val and Dim.DP in config.dims_val:
            dp_val = config.dims_val[Dim.DP]
            mbs_val = config.dims_val[Dim.MBS]
            if isinstance(dp_val, int) and isinstance(mbs_val, int) and dp_val > 0 and mbs_val > 0:
                mbn = self.global_batch_size // dp_val // mbs_val
                config.dims_val[Dim.MBN] = mbn
                injected_mbn = True
        self.config.set_parallel_config(config)
        if injected_mbn:
            del config.dims_val[Dim.MBN]
        peak_mem = self.memory_estim()
        est_kwargs = {
            "debugger": debugger,
            "device_type": self.machine.device,

        dim_coeffs = fallback.flop_coeffs.get(dimension, {})
        result = dim_coeffs.get(sub_key, {})
        if result:
            return result
    return {}


def _compute_hsdp_features(cfg, d_shard_val, device_type, mb=1):
    """Compute shared derived features for HSDP/FSDP FLOP-mode models.

    Pipeline * m cancels /mb, yielding total. Rank=rank(total).
    """
    f = _compute_hsdp_features(cfg, d_shard_val, device_type, mb=mb)
    if f["d_replicate"] <= 1:
        return 0.0
    c = _get_flop_coeffs(device_type, "comp", "hsdp")
    if not c:
        return 0.0
    total = (
        c.get("INTERCEPT", 0)
        + c.get("TP", 0) * f["tp"]
        + c.get("INV_TP", 0) * f["inv_tp"]

    Returns:
        Effective bandwidth in bytes/s.
    """
    msg_mb = msg_bytes / 1e6
    if msg_mb <= 0:
        return peak_bw_gbps * 1e9 * small_eff
    t = max(0.0, min(1.0, (log10(max(msg_mb, 1e-6)) - _LOG10_MB_FLOOR) / _LOG10_MB_RANGE))
    eff = small_eff + (large_eff - small_eff) * t
    return peak_bw_gbps * 1e9 * eff


def _cp_resolve_topology(cp, device_per_node, bw_intra, bw_inter):
    """Resolve CP topology and effective bandwidth.

        if cfg.n_exp > 1:
            fsdp_intra_vol += exp * bc * 2
        d_shard_local = cfg.d_shard_or_d
        if getattr(cfg, "comm_hsdp", 0) > 0 and d_shard_local < cfg.d:
            sharded_non_exp = non_exp / (d_shard_local * cfg.cp * cfg.t)
            sharded_exp = exp / (d_shard_local * cfg.cp * cfg.t_exp) if cfg.n_exp > 1 else 0
            hsdp_inter_vol += (sharded_non_exp + sharded_exp) * bc
    fsdp_layer_count = 1 if is_fsdp_layer else 0

    comm[Dim.TP] += EvalLayerComm.tp_comm_layer(cfg, ctx, 1)
    tp_layer_count += 1

            ag_vol = 1-1/d_shard, dp = d_rep*d_shard.
      rank=10/10, Spearman=1.0, MAPE=0.49%, 10 configs.
    """
    if fsdp_layer_count <= 0 or d_shard_val <= 0:
        return 0.0

    f = _compute_hsdp_features(cfg, d_shard_val, device_type, mb=mb)
    f["pp"] = max(pp, 1)
    is_hsdp = f["d_replicate"] > 1

    sub_key = "hsdp" if is_hsdp else "fsdp"
    c = _get_flop_coeffs(device_type, "shard", sub_key)
    if not c:
        return 0.0

    if is_hsdp:
        total = (
            c.get("INTERCEPT", 0)

    f = _compute_hsdp_features(cfg, d_shard_val, device_type, mb=mb)

    if (cfg.fsdp or d_shard_val > 1) and d_shard_val > 0:
        if f["sg"] <= 1:
            return 0.0

        if f["d_replicate"] > 1:
            c = _get_flop_coeffs(device_type, "dp", "hsdp")
            if not c:
                return 0.0
            dp_val = f["d_replicate"] * f["d_shard"]
            total = (
                c.get("INTERCEPT", 0)
                + c.get("TP", 0) * f["tp"]

                f["cross"], score, f["m"], comm_dp_raw,
            )
            return score

        return 0.0

    logger.info(
        "FLOP_DP_FALLBACK: d=%d tp=%d fsdp=%s d_shard=%d raw=%.4f result=0.0",
        f["d"], f["tp"], cfg.fsdp, f["d_shard"],

    if f["d_replicate"] <= 1:
        return 0.0
    c = _get_flop_coeffs(device_type, "tp", "hsdp")
    if not c:
        return 0.0
    total = (
        c.get("INTERCEPT", 0)
        + c.get("INV_TP", 0) * f["inv_tp"]
        + c.get("INV_SG", 0) * f["inv_sg"]

    Pipeline * m cancels /mb, yielding total. Rank=rank(total).
    """
    f = _compute_hsdp_features(cfg, d_shard_val, device_type, mb=mb)
    if f["d_replicate"] <= 1:
        return 0.0
    c = _get_flop_coeffs(device_type, "pp", "hsdp")
    if not c:
        return 0.0
    total = (
        c.get("INTERCEPT", 0)
        + c.get("AG_VOL_D_REP", 0) * f["ag_vol_d_rep"]
        + c.get("CROSS_SG", 0) * f["cross_sg"]

        comm[Dim.TP] = _flop_mode_tp_comm(
            cfg, d_shard_val, param["device_type"], mb=mb,
        )
    else:
        c = _get_flop_coeffs(param["device_type"], "tp", "fsdp")
        if c:
            tp_scaling = (c.get("A", 0) + c.get("B", 0) * f["tp"] + c.get("C", 0) / f["tp"]) / max(f["d"], 1)
            comm[Dim.TP] *= max(1, f["tp"] // dev_per_node) * tp_scaling

    comm[Dim.EP] *= max(1, param["cfg"].ep // dev_per_node)
    comm[Dim.CP] *= max(1, param["cfg"].cp // dev_per_node)

        idx = level.value - 1 if isinstance(level, NetworkLevel) else level - 1
        p2p_eff = getattr(device, 'p2p_efficiency', None)
        if p2p_eff is not None and 0 <= idx < len(p2p_eff):
            return p2p_eff[idx]
        if 0 <= idx < len(device.level_efficiency):
            return device.level_efficiency[idx]
    raise ValueError(
        f"No efficiency for level {level}; device required"
    )

        idx = level.value - 1 if isinstance(level, NetworkLevel) else level - 1
        p2p = getattr(device, 'p2p_bandwidth', None)
        if p2p is not None and idx < len(p2p):
            return p2p[idx]
        if 0 <= idx < len(device.level_bandwidth):
            return device.level_bandwidth[idx]
    raise ValueError(
        f"No P2P bandwidth for level {level}; device required"
    )

    FSDP/HSDP all-gather/reduce-scatter happens across d_shard*cp*t ranks.
    This function returns a list of n_level per hierarchy level for that
    shard group, independent of how DP/TP/CP are individually assigned.
    """
    remaining = shard_size
    levels = []
    for level in range(device.levels):
        bound = device.level_bound_number[level]
        if bound:
            n = min(remaining, bound)
            remaining = remaining // n if n > 0 else remaining
        else:
            n = remaining
        levels.append(n)
    return levels


def estimate_comm_score(
    cfg: Any,

    d_shard = cfg.d_shard_or_d

    if shard_group_size > 0:
        n_levels = _shard_group_levels(device, shard_group_size)
    else:
        assignment = device.level_assign(
            dp=cfg.d, tp=cfg.t, cp=cfg.cp, pp=cfg.p, d_shard=d_shard
        )

        n_level = n_levels[level]
        if n_level <= 1:
            continue
        if per_rank_msg:
            msg_per_rank = comm_volume
        elif shard_group_size > 0:
            prod_from_l = 1
            for k in range(level, len(n_levels)):
                prod_from_l *= n_levels[k]
            msg_per_rank = comm_volume / prod_from_l
        else:
            msg_per_rank = comm_volume / n_level
        if isinstance(a2a_efficiency, dict):
            eff = a2a_efficiency.get(level, level_efficiency(NetworkLevel(level + 1), device=device))
        elif a2a_efficiency is not None:
            eff = a2a_efficiency
        else:
            eff = level_efficiency(NetworkLevel(level + 1), device=device)
        bw_bps = device.level_bandwidth[level] * 1e9 * eff
        lat = level_latency(NetworkLevel(level + 1), device=device)

            eff = level_efficiency(NetworkLevel(level + 1), device=device)
        bw_bps = device.level_bandwidth[level] * 1e9 * eff
        lat = level_latency(NetworkLevel(level + 1), device=device)
        if per_rank_msg:
            collective_time = (
                (n_level - 1) * msg_per_rank / bw_bps
                + lat * 2 * (n_level - 1)
            )
        else: